add multibroadcast ref_ops tests

test/ref_ops_test.cpp

@@ -4181,6 +4181,71 @@ TEST_CASE(mul_dyn_test)
     EXPECT(migraphx::verify_range(results_vector, gold));
 }
 
+TEST_CASE(multibroadcast_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape a_shape{migraphx::shape::int32_type, {2, 2}};
+    std::vector<int32_t> a_data{0, 0, 0, 0};
+    migraphx::shape b_shape{migraphx::shape::int32_type, {2}};
+    std::vector<int32_t> b_data{-2, -3};
+    auto l1 = mm->add_literal(migraphx::literal{a_shape, a_data});
+    auto l2 = mm->add_literal(migraphx::literal{b_shape, b_data});
+    mm->add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", l1->get_shape().lens()}}),
+                        l2);
+    p.compile(migraphx::ref::target{});
+    auto result = p.eval({}).back();
+    auto output = result.get<int32_t>();
+    EXPECT(output(0, 0) == -2);
+    EXPECT(output(0, 1) == -3);
+    EXPECT(output(1, 0) == -2);
+    EXPECT(output(1, 1) == -3);
+}
+
+TEST_CASE(multibroadcast_2in_static_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape a_shape{migraphx::shape::int32_type, {2, 2}};
+    std::vector<int32_t> a_data{0, 0, 0, 0};
+    migraphx::shape b_shape{migraphx::shape::int32_type, {2}};
+    std::vector<int32_t> b_data{-2, -3};
+    auto l1 = mm->add_literal(migraphx::literal{a_shape, a_data});
+    auto l2 = mm->add_literal(migraphx::literal{b_shape, b_data});
+    mm->add_instruction(migraphx::make_op("multibroadcast"), l2, l1);
+    p.compile(migraphx::ref::target{});
+    auto result = p.eval({}).back();
+    auto output = result.get<int32_t>();
+    EXPECT(output(0, 0) == -2);
+    EXPECT(output(0, 1) == -3);
+    EXPECT(output(1, 0) == -2);
+    EXPECT(output(1, 1) == -3);
+}
+
+TEST_CASE(multibroadcast_2in_dyn_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape a_shape{migraphx::shape::int32_type, {{2, 4, 0}, {2, 2, 0}}};
+    migraphx::shape b_shape{migraphx::shape::int32_type, {2}};
+    std::vector<int32_t> b_data{-2, -3};
+    auto l1 = mm->add_parameter("a", a_shape);
+    auto l2 = mm->add_literal(migraphx::literal{b_shape, b_data});
+    mm->add_instruction(migraphx::make_op("multibroadcast"), l2, l1);
+    p.compile(migraphx::ref::target{});
+
+    std::vector<int32_t> a_data{0, 0, 0, 0};
+    migraphx::parameter_map params0;
+    migraphx::shape input_fixed_shape0{migraphx::shape::float_type, {2, 2}};
+    params0["a"] = migraphx::argument(input_fixed_shape0, a_data.data());
+    auto result  = p.eval(params0).back();
+    auto output  = result.get<int32_t>();
+    EXPECT(output(0, 0) == -2);
+    EXPECT(output(0, 1) == -3);
+    EXPECT(output(1, 0) == -2);
+    EXPECT(output(1, 1) == -3);
+}
+
 TEST_CASE(multinomial_test)
 {
     migraphx::program p;