Add Qlinearconcat op (#2476)

src/onnx/parse_qlinearconcat.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2023 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.
+ */
+
+#include <migraphx/onnx/op_parser.hpp>
+#include <migraphx/onnx/padding.hpp>
+#include <migraphx/onnx/conv.hpp>
+#include <migraphx/ranges.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/onnx/checks.hpp>
+#include <migraphx/onnx/broadcast_qdq.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/stringutils.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace onnx {
+
+struct parse_qlinearconcat : op_parser<parse_qlinearconcat>
+{
+    std::vector<op_desc> operators() const { return {{"QLinearConcat"}}; }
+
+    // basic type checking for QLinearConcat Operator
+    void check_inputs(const std::vector<instruction_ref>& args) const
+    {
+        auto args_size = args.size();
+        // at least 5 input tensors:
+        // 1. is Y_scale: tensor(float)
+        // 2. is Y_zero_pont: tensor(uint8)/tensor(int8)
+        // remaining is a sequence of :
+        //     3. Tensor: tensor(uint8)/tensor(int8)
+        //     4. Scale: tensor(float),
+        //     5. ZeroPoint: tensor(uint8)/tensor(int8) tensors
+        // Size can be 5, 8, 11 ...
+        if((args_size < 5) or ((args_size - 2) % 3 != 0))
+            MIGRAPHX_THROW("QLINEARCONCAT: missing inputs");
+
+        auto y_zp      = args[1];
+        auto y_zp_type = y_zp->get_shape().type();
+        if(y_zp_type != migraphx::shape::int8_type and y_zp_type != migraphx::shape::uint8_type)
+            MIGRAPHX_THROW("QLINEARCONCAT: unsupported output type");
+
+        auto t0_type = args[2]->get_shape().type();
+        if(t0_type != migraphx::shape::int8_type and t0_type != migraphx::shape::uint8_type)
+            MIGRAPHX_THROW("QLINEARCONCAT: unsupported input type");
+        for(auto idx = 2; idx < args.size(); idx += 3)
+        {
+            if((args[idx]->get_shape().type() != t0_type) or
+               (args[idx + 2]->get_shape().type() != t0_type))
+            {
+                MIGRAPHX_THROW("QLINEARCONCAT: mismatching input types");
+            }
+        }
+    }
+
+    instruction_ref parse(const op_desc& /* opd */,
+                          const onnx_parser& parser,
+                          const onnx_parser::node_info& info,
+                          const std::vector<instruction_ref>& args) const
+    {
+        check_inputs(args);
+        if(not contains(info.attributes, "axis"))
+            MIGRAPHX_THROW("QLINEARCONCAT: missing axis attribute");
+
+        auto axis = parser.parse_value(info.attributes.at("axis")).template at<int64_t>();
+        std::vector<instruction_ref> tmp;
+        for(auto idx = 2; idx < args.size(); idx += 3)
+        {
+            auto data_tensor = args[idx];
+            auto scale       = args[idx + 1];
+            auto zero_pt     = args[idx + 2];
+            tmp.push_back(bcast_qdq_instr("dequantizelinear", data_tensor, scale, zero_pt, info));
+        }
+        auto y = info.add_instruction(migraphx::make_op("concat", {{"axis", axis}}), tmp);
+
+        auto y_scale   = args[0];
+        auto y_zero_pt = args[1];
+
+        return bcast_qdq_instr("quantizelinear", y, y_scale, y_zero_pt, info);
+    }
+};
+
+} // namespace onnx
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

test/onnx/gen_onnx.py

@@ -6251,6 +6251,56 @@ def qlinearaveragepool_nt_cip_test():
     return ([node], [x], [y], [x_scale, x_zero_point, y_scale, y_zero_point])
 
 
+@onnx_test()
+def qlinearconcat_test():
+    y_scale = helper.make_tensor('1', TensorProto.FLOAT, [], [0.5])
+    y_zero_point = helper.make_tensor('2', TensorProto.INT8, [], [2])
+
+    t0 = helper.make_tensor_value_info('t0', TensorProto.INT8, [2])
+    s0 = helper.make_tensor('3', TensorProto.FLOAT, [], [0.5])
+    zp0 = helper.make_tensor('4', TensorProto.INT8, [], [1])
+
+    t1 = helper.make_tensor_value_info('t1', TensorProto.INT8, [3])
+    s1 = helper.make_tensor('5', TensorProto.FLOAT, [], [0.25])
+    zp1 = helper.make_tensor('6', TensorProto.INT8, [], [0])
+
+    y = helper.make_tensor_value_info('out', TensorProto.INT8, [5])
+
+    node = onnx.helper.make_node(
+        'QLinearConcat',
+        inputs=['1', '2', 't0', '3', '4', 't1', '5', '6'],
+        axis=0,
+        outputs=['out'],
+    )
+
+    return ([node], [t0, t1], [y], [y_scale, y_zero_point, s0, zp0, s1, zp1])
+
+
+@onnx_test()
+def qlinearconcat_3d_test():
+    y_scale = helper.make_tensor('1', TensorProto.FLOAT, [], [0.5])
+    y_zero_point = helper.make_tensor('2', TensorProto.INT8, [], [2])
+
+    t0 = helper.make_tensor_value_info('t0', TensorProto.INT8, [3, 4, 2])
+    s0 = helper.make_tensor('3', TensorProto.FLOAT, [], [0.5])
+    zp0 = helper.make_tensor('4', TensorProto.INT8, [], [10])
+
+    t1 = helper.make_tensor_value_info('t1', TensorProto.INT8, [3, 2, 2])
+    s1 = helper.make_tensor('5', TensorProto.FLOAT, [], [0.4])
+    zp1 = helper.make_tensor('6', TensorProto.INT8, [], [20])
+
+    y = helper.make_tensor_value_info('out', TensorProto.UINT8, [3, 6, 2])
+
+    node = onnx.helper.make_node(
+        'QLinearConcat',
+        inputs=['1', '2', 't0', '3', '4', 't1', '5', '6'],
+        axis=1,
+        outputs=['out'],
+    )
+
+    return ([node], [t0, t1], [y], [y_scale, y_zero_point, s0, zp0, s1, zp1])
+
+
 @onnx_test()
 def qlinearconv_test():
     # https://xadupre.github.io/draft/onnx/onnx_doc_folder/onnx__QLinearConv.html

test/onnx/onnx_test.cpp

@@ -5645,6 +5645,59 @@ TEST_CASE(qlinearaveragepool_notset_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(qlinearconcat_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+
+    auto sc_y   = mm->add_literal(migraphx::literal{migraphx::shape::float_type, {0.5}});
+    auto z_pt_y = mm->add_literal(migraphx::literal{migraphx::shape::int8_type, {2}});
+
+    auto sc_0   = mm->add_literal(migraphx::literal{migraphx::shape::float_type, {0.5}});
+    auto z_pt_0 = mm->add_literal(migraphx::literal{migraphx::shape::int8_type, {1}});
+
+    auto sc_1   = mm->add_literal(migraphx::literal{migraphx::shape::float_type, {0.25}});
+    auto z_pt_1 = mm->add_literal(migraphx::literal{migraphx::shape::int8_type, {0}});
+
+    auto t0 = mm->add_parameter("t0", {migraphx::shape::int8_type, {2}});
+    auto t1 = mm->add_parameter("t1", {migraphx::shape::int8_type, {3}});
+
+    auto scale_0_bcast =
+        mm->add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {2}}}), sc_0);
+
+    auto z_pt_0_bcast =
+        mm->add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {2}}}), z_pt_0);
+
+    auto fp_0 =
+        mm->add_instruction(migraphx::make_op("dequantizelinear"), t0, scale_0_bcast, z_pt_0_bcast);
+
+    auto scale_1_bcast =
+        mm->add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {3}}}), sc_1);
+
+    auto z_pt_1_bcast =
+        mm->add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {3}}}), z_pt_1);
+
+    auto fp_1 =
+        mm->add_instruction(migraphx::make_op("dequantizelinear"), t1, scale_1_bcast, z_pt_1_bcast);
+
+    auto fp_y = mm->add_instruction(migraphx::make_op("concat", {{"axis", 0}}), fp_0, fp_1);
+
+    auto scale_y_bcast =
+        mm->add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {5}}}), sc_y);
+
+    auto z_pt_y_bcast =
+        mm->add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {5}}}), z_pt_y);
+
+    auto y =
+        mm->add_instruction(migraphx::make_op("quantizelinear"), fp_y, scale_y_bcast, z_pt_y_bcast);
+
+    mm->add_return({y});
+
+    auto prog = migraphx::parse_onnx("qlinearconcat_test.onnx");
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(qlinearconv_test)
 {
     migraphx::program p;

test/onnx/verify_onnx.cpp

@@ -1932,6 +1932,52 @@ TEST_CASE(qlinearaveragepool_nt_cip_test)
     EXPECT(migraphx::verify::verify_rms_range(result_vector, gold));
 }
 
+TEST_CASE(qlinearconcat_test)
+{
+    auto p = migraphx::parse_onnx("qlinearconcat_test.onnx");
+    p.compile(migraphx::make_target("ref"));
+
+    std::vector<int8_t> data_t0 = {2, 3};
+    migraphx::shape s_t0{migraphx::shape::int8_type, {2}};
+    migraphx::parameter_map pp;
+    pp["t0"] = migraphx::argument(s_t0, data_t0.data());
+
+    std::vector<int8_t> data_t1 = {6, 8, 10};
+    migraphx::shape s_t1{migraphx::shape::int8_type, {3}};
+    pp["t1"] = migraphx::argument(s_t1, data_t1.data());
+
+    auto result = p.eval(pp).back();
+    std::vector<int8_t> result_vector;
+    result.visit([&](auto output) { result_vector.assign(output.begin(), output.end()); });
+
+    std::vector<int8_t> gold = {3, 4, 5, 6, 7};
+    EXPECT(migraphx::verify::verify_rms_range(result_vector, gold));
+}
+
+TEST_CASE(qlinearconcat_3d_test)
+{
+    auto p = migraphx::parse_onnx("qlinearconcat_3d_test.onnx");
+    p.compile(migraphx::make_target("ref"));
+
+    std::vector<int8_t> data_t0 = {10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
+                                   10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10};
+    migraphx::shape s_t0{migraphx::shape::int8_type, {3, 4, 2}};
+    migraphx::parameter_map pp;
+    pp["t0"] = migraphx::argument(s_t0, data_t0.data());
+
+    std::vector<int8_t> data_t1 = {25, 25, 25, 25, 25, 25, 25, 25, 25, 25, 25, 25};
+    migraphx::shape s_t1{migraphx::shape::int8_type, {3, 2, 2}};
+    pp["t1"] = migraphx::argument(s_t1, data_t1.data());
+
+    auto result = p.eval(pp).back();
+    std::vector<uint8_t> result_vector;
+    result.visit([&](auto output) { result_vector.assign(output.begin(), output.end()); });
+
+    std::vector<int8_t> gold = {2, 2, 2, 2, 2, 2, 2, 2, 6, 6, 6, 6, 2, 2, 2, 2, 2, 2,
+                                2, 2, 6, 6, 6, 6, 2, 2, 2, 2, 2, 2, 2, 2, 6, 6, 6, 6};
+    EXPECT(migraphx::verify::verify_rms_range(result_vector, gold));
+}
+
 TEST_CASE(qlinearconv_test)
 {
     // https://xadupre.github.io/draft/onnx/onnx_doc_folder/onnx__QLinearConv.html