    def _test_vectorization(
        self,
        batch_sizes,
        ms,
        ks,
        test_name,
        expected_max_read_t,
        expected_read_types,
        expected_op_t,
        expected_data_t,
        ns=None,
        dtype="float16",
    ):
        """
        Test add(add(X0(B, M0, K0, N0), X1(B, M1, K1, N1)), X2(B, M2, K2, N2))
        """

        batch_dim = shape_utils.gen_int_var_min_max(batch_sizes, name="batch_dim")
        if ns is None:
            ns = [1, 1, 1]

        X0 = Tensor(
            shape=[batch_dim, IntImm(ms[0]), IntImm(ks[0]), IntImm(ns[0])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X1 = Tensor(
            shape=[batch_dim, IntImm(ms[1]), IntImm(ks[1]), IntImm(ns[1])],
            dtype=dtype,
            name="input1",
            is_input=True,
        )
        X2 = Tensor(
            shape=[batch_dim, IntImm(ms[2]), IntImm(ks[2]), IntImm(ns[2])],
            dtype=dtype,
            name="input2",
            is_input=True,
        )
        add_1 = ops.elementwise(FuncEnum.ADD)(X0, X1)
        output = ops.elementwise(FuncEnum.ADD)(add_1, X2)
        output._attrs["name"] = "output0"
        output._attrs["is_output"] = True

        target = detect_target()
        module = compile_model(output, target, "./tmp", test_name)

        debug_sorted_graph = module.debug_sorted_graph
        sorted_ops = graph_utils.get_sorted_ops(debug_sorted_graph)
        self.assertEqual(len(sorted_ops), 1)
        self.assertEqual(sorted_ops[0]._attrs["op"], "fused_elementwise")
        self.assertEqual(sorted_ops[0]._attrs["max_read_t"], expected_max_read_t)
        self.assertEqual(
            self._get_sorted_read_types(sorted_ops[0]), expected_read_types
        )
        self.assertEqual(sorted_ops[0]._attrs["op_t"], expected_op_t)
        self.assertEqual(sorted_ops[0]._attrs["data_t"], expected_data_t)

        for batch_size in batch_sizes:
            x0_pt = get_random_torch_tensor(
                [batch_size, ms[0], ks[0], ns[0]], dtype=dtype
            )
            x1_pt = get_random_torch_tensor(
                [batch_size, ms[1], ks[1], ns[1]], dtype=dtype
            )
            x2_pt = get_random_torch_tensor(
                [batch_size, ms[2], ks[2], ns[2]], dtype=dtype
            )
            output_pt = (x0_pt + x1_pt) + x2_pt
            inputs = {"input0": x0_pt, "input1": x1_pt, "input2": x2_pt}
            output = torch.empty_like(output_pt)
            module.run_with_tensors(inputs, [output])
            self.assertTrue(torch.allclose(output, output_pt, atol=1e-2, rtol=1e-2))
Test add(add(X0(B, M0, K0, N0), X1(B, M1, K1, N1)), X2(B, M2, K2, N2))
tests/unittest/ops/test_fused_elementwise_broadcast.py