| 36 | inline namespace MIGRAPHX_INLINE_NS { |
| 37 | |
| 38 | static void int4_quantize_module(module& m) |
| 39 | { |
| 40 | std::vector<std::string> int4_instrs{"dot", "convolution"}; |
| 41 | |
| 42 | for(auto ins : iterator_for(m)) |
| 43 | { |
| 44 | if(not(contains(int4_instrs, ins->name()))) |
| 45 | continue; |
| 46 | |
| 47 | if(ins->inputs().empty()) |
| 48 | continue; |
| 49 | |
| 50 | auto s = ins->get_shape(); |
| 51 | |
| 52 | auto mod_inputs = ins->module_inputs(); |
| 53 | |
| 54 | // Convert each of the inputs that are fp32 or fp16 to int4 |
| 55 | auto inputs = ins->inputs(); |
| 56 | std::transform(inputs.begin(), inputs.end(), inputs.begin(), [&](auto inp) { |
| 57 | auto sh = inp->get_shape(); |
| 58 | if(sh.broadcasted()) |
| 59 | return inp; |
| 60 | auto input_type = sh.type(); |
| 61 | if(input_type != shape::float_type and input_type != shape::half_type) |
| 62 | return inp; |
| 63 | auto lens = sh.lens(); |
| 64 | if(lens[lens.size() - 1] % 2) |
| 65 | return inp; // even sized dimensions to pack |
| 66 | |
| 67 | if(not inp->can_eval()) |
| 68 | return inp; |
| 69 | |
| 70 | std::vector<float> val; |
| 71 | inp->eval().visit([&](auto in_data) { val.assign(in_data.begin(), in_data.end()); }); |
| 72 | |
| 73 | auto [min, max] = std::minmax_element(val.begin(), val.end()); |
| 74 | *min = *min > 0 ? 0 : *min; |
| 75 | *max = *max < 0 ? 0 : *max; |
| 76 | float fscale4 = (*max - *min) / 15; // INT4 range is [0-15] |
| 77 | int zp4 = float_equal(fscale4, 0) ? 0 : std::round(-*min / fscale4); |
| 78 | |
| 79 | auto scale = m.add_literal(literal({s.type()}, {fscale4})); |
| 80 | scale = |
| 81 | m.insert_instruction(ins, make_op("multibroadcast", {{"out_lens", lens}}), scale); |
| 82 | auto zp = m.add_literal(literal{{shape::uint8_type}, {zp4}}); |
| 83 | zp = m.insert_instruction(ins, make_op("multibroadcast", {{"out_lens", lens}}), zp); |
| 84 | auto q_in = m.insert_instruction(ins, make_op("quantizelinear"), inp, scale, zp); |
| 85 | |
| 86 | auto pk = m.insert_instruction(ins, make_op("pack_int4", {{"axis", -1}}), q_in); |
| 87 | auto unpk = m.insert_instruction(ins, make_op("unpack_int4", {{"axis", -1}}), pk); |
| 88 | |
| 89 | auto dq_scale = m.add_literal(literal({s.type()}, {fscale4})); |
| 90 | dq_scale = m.insert_instruction( |
| 91 | ins, make_op("multibroadcast", {{"out_lens", lens}}), dq_scale); |
| 92 | |
| 93 | auto dq_zp = m.add_literal(literal{{shape::uint8_type}, {zp4}}); |
| 94 | dq_zp = |
| 95 | m.insert_instruction(ins, make_op("multibroadcast", {{"out_lens", lens}}), dq_zp); |
no test coverage detected