simplify_qdq.cpp 9.92 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
/*
 * The MIT License (MIT)
 *
 * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
turneram's avatar
turneram committed
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
#include <migraphx/simplify_qdq.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/iterator_for.hpp>
#include <migraphx/make_op.hpp>
#include <migraphx/program.hpp>
#include <migraphx/shape.hpp>
#include <migraphx/matcher.hpp>
#include <migraphx/dead_code_elimination.hpp>
#include <migraphx/pass_manager.hpp>
#include <migraphx/op/convolution.hpp>
#include <migraphx/op/quant_convolution.hpp>
#include <migraphx/op/dot.hpp>
#include <migraphx/op/quant_dot.hpp>
#include <migraphx/register_op.hpp>

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {

std::unordered_set<std::string> get_quantizable_op_names()
{
    static std::unordered_set<std::string> s = {"convolution", "dot"};
    return s;
}

48
struct match_find_quantizable_ops
turneram's avatar
turneram committed
49
{
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
    static bool
    is_valid_scale(instruction_ref scale, std::vector<std::size_t> lens, std::size_t axis)
    {
        return scale->get_shape().scalar() or scale->get_shape().elements() == lens.at(axis);
    }

    static bool is_valid_zero_point(instruction_ref zp)
    {
        if(not zp->can_eval())
            return false;

        bool all_zeros = false;
        zp->eval().visit([&](auto z) {
            all_zeros =
                std::all_of(z.begin(), z.end(), [&](auto val) { return float_equal(val, 0); });
turneram's avatar
turneram committed
65
        });
66
67
        return all_zeros;
    }
turneram's avatar
turneram committed
68

69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
    static auto
    scale_broadcast_op(instruction_ref scale, std::vector<std::size_t> lens, std::size_t axis)
    {
        if(scale->get_shape().scalar())
        {
            return migraphx::make_op("multibroadcast", {{"out_lens", lens}});
        }
        else
        {
            return migraphx::make_op("broadcast", {{"out_lens", lens}, {"axis", axis}});
        }
    }

    // Helper function to insert quantized versions of any broadcasts and transpose ops that
    // occur between dequantizelinear and the quantized op
    static auto
85
    propagate_quantized_ins(module& m, const instruction_ref dqins, const instruction_ref qop_arg)
86
    {
87
88
89
90
91
92
93
94
95
96
97
        auto prev_ins = qop_arg;
        std::vector<instruction_ref> ins_inbetween;
        // matcher skips continguous, multi/broadcasts and transposes, collect all those
        // instructions
        while(prev_ins != dqins)
        {
            ins_inbetween.push_back(prev_ins);
            prev_ins = prev_ins->inputs().front();
        }
        auto qinp = dqins->inputs().front();
        for(auto ins : reverse_iterator_for(ins_inbetween))
98
        {
99
            qinp = m.insert_instruction(dqins, (*ins)->get_operator(), {qinp});
100
101
102
        }
        return qinp;
    }
turneram's avatar
turneram committed
103

104
    static auto dequantizelinear_op(const std::string& scale, const std::string& zp)
turneram's avatar
turneram committed
105
106
    {
        return match::name("dequantizelinear")(
107
108
109
            match::arg(0)(match::skip(match::name("quantizelinear"))(match::any())),
            match::arg(1)(match::skip_broadcasts(match::is_constant().bind(scale))),
            match::arg(2)(match::skip_broadcasts(match::is_constant().bind(zp))));
turneram's avatar
turneram committed
110
111
112
113
114
    }

    auto matcher() const
    {
        return match::name(get_quantizable_op_names())(
115
116
117
118
            match::arg(0)(match::skip_broadcasts_transposes_contiguous(
                dequantizelinear_op("scale1", "zp1").bind("dq1"))),
            match::arg(1)(match::skip_broadcasts_transposes_contiguous(
                dequantizelinear_op("scale2", "zp2").bind("dq2"))));
turneram's avatar
turneram committed
119
120
    }

121
    void apply(module& m, const match::matcher_result& r) const
turneram's avatar
turneram committed
122
123
    {
        auto qop    = r.result;
124
125
        auto dq1    = r.instructions["dq1"];
        auto dq2    = r.instructions["dq2"];
turneram's avatar
turneram committed
126
127
        auto scale1 = r.instructions["scale1"];
        auto scale2 = r.instructions["scale2"];
128
129
        auto zp1    = r.instructions["zp1"];
        auto zp2    = r.instructions["zp2"];
130
131
132
133
134
        // Only INT8 or FP8 type currently supported
        std::set<migraphx::shape::type_t> supported_types = {migraphx::shape::fp8e4m3fnuz_type,
                                                             migraphx::shape::int8_type};
        if(not contains(supported_types, dq1->inputs().front()->get_shape().type()) or
           not contains(supported_types, dq2->inputs().front()->get_shape().type()))
135
136
137
138
            return;

        // Only symmetric quantization supported (ie. non-zero zero_points not allowed)
        if(not(is_valid_zero_point(zp1) and is_valid_zero_point(zp2)))
turneram's avatar
turneram committed
139
140
            return;

141
142
143
        // Only support scalar and 1D scales
        if(scale1->get_shape().lens().size() != 1 or scale2->get_shape().lens().size() != 1)
            return;
turneram's avatar
turneram committed
144

145
        // Propagate q1 and q2 through any broadcasts and transposes before qop
turneram's avatar
turneram committed
146
        auto qop_args  = qop->inputs();
147
148
        qop_args.at(0) = propagate_quantized_ins(m, dq1, qop_args[0]);
        qop_args.at(1) = propagate_quantized_ins(m, dq2, qop_args[1]);
turneram's avatar
turneram committed
149
        instruction_ref dq;
150
        instruction_ref out_scale;
turneram's avatar
turneram committed
151
152
153
154
155
156
        instruction_ref zero_point;
        if(qop->name() == "convolution")
        {
            auto conv_val = qop->get_operator().to_value();
            dq            = m.insert_instruction(
                qop, migraphx::make_op("quant_convolution", conv_val), qop_args);
157
158
159
160
161
162
163
164
165
166
167
168
169
            auto out_lens = dq->get_shape().lens();

            // Input scale should always be scalar and weight scale can be scalar or 1D of the
            // same lens as the output channel dim (dim 1 in the output)
            if(not(is_valid_scale(scale1, out_lens, 1) and is_valid_scale(scale2, out_lens, 1)))
                return;

            auto s1_bcast =
                m.insert_instruction(qop, scale_broadcast_op(scale1, out_lens, 1), scale1);
            auto s2_bcast =
                m.insert_instruction(qop, scale_broadcast_op(scale2, out_lens, 1), scale2);

            out_scale = m.insert_instruction(qop, migraphx::make_op("mul"), s1_bcast, s2_bcast);
turneram's avatar
turneram committed
170
171
172
        }
        else if(qop->name() == "dot")
        {
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
            dq            = m.insert_instruction(qop, migraphx::make_op("quant_dot"), qop_args);
            auto out_lens = dq->get_shape().lens();

            // For (..., M, N) x (..., N, K) dot, only support cases where quantization axis is M
            // for input1 and K for input 2
            if(not(is_valid_scale(scale1, out_lens, out_lens.size() - 2) and
                   is_valid_scale(scale2, out_lens, out_lens.size() - 1)))
                return;

            auto s1_bcast = m.insert_instruction(
                qop, scale_broadcast_op(scale1, out_lens, out_lens.size() - 2), scale1);
            auto s2_bcast = m.insert_instruction(
                qop, scale_broadcast_op(scale2, out_lens, out_lens.size() - 1), scale2);

            out_scale = m.insert_instruction(qop, migraphx::make_op("mul"), s1_bcast, s2_bcast);
turneram's avatar
turneram committed
188
189
        }

190
        dq = m.insert_instruction(qop, make_op("dequantizelinear"), dq, out_scale);
turneram's avatar
turneram committed
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
        m.replace_instruction(qop, dq);
    }
};

bool compare_literals(instruction_ref ins1, instruction_ref ins2)
{
    if(ins1->name() == "broadcast" or ins1->name() == "multibroadcast")
        ins1 = ins1->inputs().front();
    auto x = ins1->eval();
    if(x.empty())
        return false;
    auto literal1 = ins1->get_literal();
    if(ins2->name() == "broadcast" or ins2->name() == "multibroadcast")
        ins2 = ins2->inputs().front();
    auto y = ins2->eval();
    if(y.empty())
        return false;
    auto literal2 = ins2->get_literal();

    bool diff_shapes_equal_vals = false;
    visit_all(ins1->get_literal(), ins2->get_literal())([&](const auto l1, const auto l2) {
        diff_shapes_equal_vals =
213
214
215
216
217
218
219
220
221
            std::all_of(l1.begin() + 1,
                        l1.end(),
                        [&](auto v) {
                            return ((float_equal(v, l1.front())) or
                                    (std::isinf(l1.front()) and std::isinf(v)));
                        }) and
            std::all_of(l2.begin(), l2.end(), [&](auto v) {
                return ((float_equal(v, l1.front())) or (std::isinf(l1.front()) and std::isinf(v)));
            });
turneram's avatar
turneram committed
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
247
248
249
250
251
252
253
254
255
256
257
    });

    return (x == y) or diff_shapes_equal_vals;
}

void remove_qdq_pairs(module& m)
{
    for(auto ins : iterator_for(m))
    {
        auto args = ins->inputs();
        for(auto&& arg : args)
        {
            if(arg->name() == "dequantizelinear")
            {
                auto q = arg->inputs().front();
                if((q->name() == "quantizelinear") and
                   compare_literals(arg->inputs().at(1), q->inputs().at(1)) and
                   compare_literals(arg->inputs().at(2), q->inputs().at(2)))
                {
                    instruction::replace_argument(ins, arg, q->inputs().front());
                }
            }
        }
    }
}

void simplify_qdq::apply(module& m) const
{
    match::find_matches(m, match_find_quantizable_ops{});
    migraphx::run_passes(m, {migraphx::dead_code_elimination{}});
    remove_qdq_pairs(m);
    migraphx::run_passes(m, {migraphx::dead_code_elimination{}});
}

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx