# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: E731, F401, F841
import tvm.testing
from tvm import relax, tirx
from tvm.relax.transform import CombineParallelMatmul
from tvm.script import relax as R
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import relax as relax_builder


def get_parallel_matmul(
    num_branches,
    lhs_shape=(640, 640),
    rhs_shape=(640, 640),
    with_bias=None,
    activation=None,
):
    dtype = "float32"

    activation_map = {"relu": R.nn.relu, "gelu": R.nn.gelu}

    with IRBuilder() as builder:
        with relax_builder.function():
            R.func_name("main")
            x = R.arg("x", R.Tensor(lhs_shape, dtype))

            rhs = []
            bias = []

            for i in range(num_branches):
                rhs.append(R.arg("y", R.Tensor(rhs_shape, dtype)))

                if with_bias and with_bias[i]:
                    bias.append(R.arg("bias", R.Tensor((rhs_shape[1],), dtype)))
                else:
                    bias.append(None)

            with R.dataflow() as frame:
                branches = []

                for i, r in enumerate(rhs):
                    result = R.emit(R.matmul(x, r, out_dtype=dtype))
                    if bias[i]:
                        result = R.emit(result + bias[i])
                    if activation and activation[i]:
                        result = R.emit(activation_map[activation[i]](result))

                    branches.append(result)

                R.output(R.emit(R.concat(branches, axis=1)))

            R.func_ret_value(frame.output_vars[0])

    func = builder.get()
    return tvm.IRModule({"main": func})


def test_simple():
    mod_orig = get_parallel_matmul(1)
    mod = CombineParallelMatmul()(mod_orig)

    tvm.ir.assert_structural_equal(mod, mod_orig)

    mod = get_parallel_matmul(3)
    mod = CombineParallelMatmul()(mod)

    @R.function
    def expected1(
        x: R.Tensor((640, 640), dtype="float32"),
        y: R.Tensor((640, 640), dtype="float32"),
        y_1: R.Tensor((640, 640), dtype="float32"),
        y_2: R.Tensor((640, 640), dtype="float32"),
    ) -> R.Tensor((640, 1920), dtype="float32"):
        with R.dataflow():
            lv = R.concat((y, y_1, y_2), axis=1)
            lv1 = R.matmul(x, lv, out_dtype="float32")
            lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=1)
            lv_1 = lv2[0]
            lv1_1 = lv2[1]
            lv2_1 = lv2[2]
            lv3 = R.concat((lv_1, lv1_1, lv2_1), axis=1)
            R.output(lv3)
        return lv3

    tvm.ir.assert_structural_equal(mod["main"], expected1.with_attr("global_symbol", "main"))

    # Test a batched LHS case, slicing is done on the axis 2
    mod = get_parallel_matmul(3, lhs_shape=(2, 1024, 640))
    mod = CombineParallelMatmul()(mod)

    @R.function
    def expected2(
        x: R.Tensor((2, 1024, 640), dtype="float32"),
        y: R.Tensor((640, 640), dtype="float32"),
        y_1: R.Tensor((640, 640), dtype="float32"),
        y_2: R.Tensor((640, 640), dtype="float32"),
    ) -> R.Tensor((2, 3072, 640), dtype="float32"):
        with R.dataflow():
            lv = R.concat((y, y_1, y_2), axis=1)
            lv1 = R.matmul(x, lv, out_dtype="float32")
            lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=2)
            lv_1 = lv2[0]
            lv1_1 = lv2[1]
            lv2_1 = lv2[2]
            lv3 = R.concat((lv_1, lv1_1, lv2_1), axis=1)
            R.output(lv3)
        return lv3

    tvm.ir.assert_structural_equal(mod["main"], expected2.with_attr("global_symbol", "main"))


def test_bias():
    mod = get_parallel_matmul(3, with_bias=[True, True, True])
    mod = CombineParallelMatmul()(mod)

    @R.function
    def expected1(
        x: R.Tensor((640, 640), dtype="float32"),
        y: R.Tensor((640, 640), dtype="float32"),
        bias: R.Tensor((640,), dtype="float32"),
        y_1: R.Tensor((640, 640), dtype="float32"),
        bias_1: R.Tensor((640,), dtype="float32"),
        y_2: R.Tensor((640, 640), dtype="float32"),
        bias_2: R.Tensor((640,), dtype="float32"),
    ) -> R.Tensor((640, 1920), dtype="float32"):
        with R.dataflow():
            lv = R.concat((y, y_1, y_2), axis=1)
            lv1 = R.matmul(x, lv, out_dtype="float32")
            lv2 = R.concat((bias, bias_1, bias_2), axis=0)
            lv3 = R.add(lv1, lv2)
            lv4 = R.split(lv3, indices_or_sections=[640, 1280], axis=1)
            lv1_1 = lv4[0]
            lv3_1 = lv4[1]
            lv5 = lv4[2]
            lv6 = R.concat((lv1_1, lv3_1, lv5), axis=1)
            R.output(lv6)
        return lv6

    tvm.ir.assert_structural_equal(mod["main"], expected1.with_attr("global_symbol", "main"))

    mod = get_parallel_matmul(3, with_bias=[True, False, True])
    mod = CombineParallelMatmul()(mod)

    @R.function
    def expected2(
        x: R.Tensor((640, 640), dtype="float32"),
        y: R.Tensor((640, 640), dtype="float32"),
        bias: R.Tensor((640,), dtype="float32"),
        y_1: R.Tensor((640, 640), dtype="float32"),
        y_2: R.Tensor((640, 640), dtype="float32"),
        bias_1: R.Tensor((640,), dtype="float32"),
    ) -> R.Tensor((640, 1920), dtype="float32"):
        with R.dataflow():
            lv = R.concat((y, y_1, y_2), axis=1)
            lv1 = R.matmul(x, lv, out_dtype="float32")
            lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=1)
            lv_1 = lv2[0]
            lv1_1 = R.add(lv_1, bias)
            lv2_1 = lv2[1]
            lv3 = lv2[2]
            lv4 = R.add(lv3, bias_1)
            lv5 = R.concat((lv1_1, lv2_1, lv4), axis=1)
            R.output(lv5)
        return lv5

    tvm.ir.assert_structural_equal(mod["main"], expected2.with_attr("global_symbol", "main"))


def test_activation():
    mod = get_parallel_matmul(3, activation=["relu", "relu", "relu"])
    mod = CombineParallelMatmul()(mod)

    @R.function
    def expected1(
        x: R.Tensor((640, 640), dtype="float32"),
        y: R.Tensor((640, 640), dtype="float32"),
        y_1: R.Tensor((640, 640), dtype="float32"),
        y_2: R.Tensor((640, 640), dtype="float32"),
    ) -> R.Tensor((640, 1920), dtype="float32"):
        with R.dataflow():
            lv = R.concat((y, y_1, y_2), axis=1)
            lv1 = R.matmul(x, lv, out_dtype="float32")
            lv2 = R.nn.relu(lv1)
            lv3 = R.split(lv2, indices_or_sections=[640, 1280], axis=1)
            lv1_1 = lv3[0]
            lv3_1 = lv3[1]
            lv5 = lv3[2]
            lv6 = R.concat((lv1_1, lv3_1, lv5), axis=1)
            R.output(lv6)
        return lv6

    tvm.ir.assert_structural_equal(mod["main"], expected1.with_attr("global_symbol", "main"))

    mod = get_parallel_matmul(3, activation=["gelu", "relu", "relu"])
    mod = CombineParallelMatmul()(mod)

    @R.function
    def expected2(
        x: R.Tensor((640, 640), dtype="float32"),
        y: R.Tensor((640, 640), dtype="float32"),
        y_1: R.Tensor((640, 640), dtype="float32"),
        y_2: R.Tensor((640, 640), dtype="float32"),
    ) -> R.Tensor((640, 1920), dtype="float32"):
        with R.dataflow():
            lv = R.concat((y, y_1, y_2), axis=1)
            lv1 = R.matmul(x, lv, out_dtype="float32")
            lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=1)
            lv_1 = lv2[0]
            lv1_1 = R.nn.gelu(lv_1)
            lv2_1 = lv2[1]
            lv3 = R.nn.relu(lv2_1)
            lv4 = lv2[2]
            lv5 = R.nn.relu(lv4)
            lv6 = R.concat((lv1_1, lv3, lv5), axis=1)
            R.output(lv6)
        return lv6

    tvm.ir.assert_structural_equal(mod["main"], expected2.with_attr("global_symbol", "main"))

    mod = get_parallel_matmul(3, activation=["relu", None, None])
    mod = CombineParallelMatmul()(mod)

    @R.function
    def expected3(
        x: R.Tensor((640, 640), dtype="float32"),
        y: R.Tensor((640, 640), dtype="float32"),
        y_1: R.Tensor((640, 640), dtype="float32"),
        y_2: R.Tensor((640, 640), dtype="float32"),
    ) -> R.Tensor((640, 1920), dtype="float32"):
        with R.dataflow():
            lv = R.concat((y, y_1, y_2), axis=1)
            lv1 = R.matmul(x, lv, out_dtype="float32")
            lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=1)

            lv_1 = lv2[0]
            lv1_1 = R.nn.relu(lv_1)
            lv2_1 = lv2[1]
            lv3 = lv2[2]
            lv4 = R.concat((lv1_1, lv2_1, lv3), axis=1)
            R.output(lv4)
        return lv4

    tvm.ir.assert_structural_equal(mod["main"], expected3.with_attr("global_symbol", "main"))


def test_bias_activation():
    mod = get_parallel_matmul(3, with_bias=[True, True, True], activation=["relu", "relu", "relu"])
    mod = CombineParallelMatmul()(mod)

    @R.function
    def expected1(
        x: R.Tensor((640, 640), dtype="float32"),
        y: R.Tensor((640, 640), dtype="float32"),
        bias: R.Tensor((640,), dtype="float32"),
        y_1: R.Tensor((640, 640), dtype="float32"),
        bias_1: R.Tensor((640,), dtype="float32"),
        y_2: R.Tensor((640, 640), dtype="float32"),
        bias_2: R.Tensor((640,), dtype="float32"),
    ) -> R.Tensor((640, 1920), dtype="float32"):
        with R.dataflow():
            lv = R.concat((y, y_1, y_2), axis=1)
            lv1 = R.matmul(x, lv, out_dtype="float32")
            lv2 = R.concat((bias, bias_1, bias_2), axis=0)
            lv3 = R.add(lv1, lv2)
            lv4 = R.nn.relu(lv3)
            lv5 = R.split(lv4, indices_or_sections=[640, 1280], axis=1)
            lv2_1 = lv5[0]
            lv5_1 = lv5[1]
            lv8 = lv5[2]
            lv9 = R.concat((lv2_1, lv5_1, lv8), axis=1)
            R.output(lv9)
        return lv9

    tvm.ir.assert_structural_equal(mod["main"], expected1.with_attr("global_symbol", "main"))

    mod = get_parallel_matmul(3, with_bias=[True, True, True], activation=["relu", None, "relu"])
    mod = CombineParallelMatmul()(mod)

    @R.function
    def expected2(
        x: R.Tensor((640, 640), dtype="float32"),
        y: R.Tensor((640, 640), dtype="float32"),
        bias: R.Tensor((640,), dtype="float32"),
        y_1: R.Tensor((640, 640), dtype="float32"),
        bias_1: R.Tensor((640,), dtype="float32"),
        y_2: R.Tensor((640, 640), dtype="float32"),
        bias_2: R.Tensor((640,), dtype="float32"),
    ) -> R.Tensor((640, 1920), dtype="float32"):
        with R.dataflow():
            lv = R.concat((y, y_1, y_2), axis=1)
            lv1 = R.matmul(x, lv, out_dtype="float32")
            lv2 = R.concat((bias, bias_1, bias_2), axis=0)
            lv3 = R.add(lv1, lv2)
            lv4 = R.split(lv3, indices_or_sections=[640, 1280], axis=1)
            lv1_1 = lv4[0]
            lv2_1 = R.nn.relu(lv1_1)
            lv4_1 = lv4[1]
            lv6 = lv4[2]
            lv7 = R.nn.relu(lv6)
            lv8 = R.concat((lv2_1, lv4_1, lv7), axis=1)
            R.output(lv8)
        return lv8

    tvm.ir.assert_structural_equal(mod["main"], expected2.with_attr("global_symbol", "main"))

    mod = get_parallel_matmul(3, with_bias=[True, False, True], activation=["relu", None, "relu"])
    mod = CombineParallelMatmul()(mod)

    @R.function
    def expected3(
        x: R.Tensor((640, 640), dtype="float32"),
        y: R.Tensor((640, 640), dtype="float32"),
        bias: R.Tensor((640,), dtype="float32"),
        y_1: R.Tensor((640, 640), dtype="float32"),
        y_2: R.Tensor((640, 640), dtype="float32"),
        bias_1: R.Tensor((640,), dtype="float32"),
    ) -> R.Tensor((640, 1920), dtype="float32"):
        with R.dataflow():
            lv = R.concat((y, y_1, y_2), axis=1)
            lv1 = R.matmul(x, lv, out_dtype="float32")
            lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=1)
            lv_1 = lv2[0]
            lv1_1 = R.add(lv_1, bias)
            lv2_1 = R.nn.relu(lv1_1)
            lv3 = lv2[1]
            lv4 = lv2[2]
            lv5 = R.add(lv4, bias_1)
            lv6 = R.nn.relu(lv5)
            lv7 = R.concat((lv2_1, lv3, lv6), axis=1)
            R.output(lv7)
        return lv7

    tvm.ir.assert_structural_equal(mod["main"], expected3.with_attr("global_symbol", "main"))


def test_rhs_batched():
    @R.function(private=True)
    def before(
        x: R.Tensor((1024, 640), "float32"),
        w0: R.Tensor((2, 640, 640), "float32"),
        w1: R.Tensor((640, 640), "float32"),
        w2: R.Tensor((2, 640, 640), "float32"),
        w3: R.Tensor((3, 4, 640, 640), "float32"),
    ):
        with R.dataflow():
            lv0 = R.matmul(x, w0)
            lv1 = R.matmul(x, w1)
            lv2 = R.matmul(x, w2)
            lv3 = R.matmul(x, w3)
            out = (lv0, lv1, lv2, lv3)
            R.output(out)
        return out

    after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]

    @R.function(private=True)
    def expected(
        x: R.Tensor((1024, 640), dtype="float32"),
        w0: R.Tensor((2, 640, 640), dtype="float32"),
        w1: R.Tensor((640, 640), dtype="float32"),
        w2: R.Tensor((2, 640, 640), dtype="float32"),
        w3: R.Tensor((3, 4, 640, 640), dtype="float32"),
    ):
        with R.dataflow():
            lv = R.concat((w0, w2), axis=2)
            lv1 = R.matmul(x, lv, out_dtype="float32")
            lv2 = R.split(lv1, indices_or_sections=[640], axis=2)
            lv0 = lv2[0]
            lv1_1 = R.matmul(x, w1, out_dtype=None)
            lv2_1 = lv2[1]
            lv3 = R.matmul(x, w3, out_dtype=None)
            out = lv0, lv1_1, lv2_1, lv3
            R.output(out)
        return out

    tvm.ir.assert_structural_equal(after, expected)

    @tvm.script.ir_module
    class four_matmul_incompatible_batches:
        @R.function
        def main(
            x: R.Tensor((1024, 640), "float32"),
            w0: R.Tensor((2, 640, 640), "float32"),
            w1: R.Tensor((3, 640, 640), "float32"),
            w2: R.Tensor((2, 640, 640), "float32"),
            w3: R.Tensor((2, 640, 640), "float32"),
        ):
            with R.dataflow():
                lv0 = R.matmul(x, w0)
                lv1 = R.matmul(x, w1)
                lv2 = R.matmul(x, w2)
                lv3 = R.matmul(x, w3)
                out = (lv0, lv1, lv2, lv3)
                R.output(out)
            return out

    mod = CombineParallelMatmul()(four_matmul_incompatible_batches)
    # For now, when rhs matrices have the same rank but different batch sizes, we don't
    # combine any of them.
    tvm.ir.assert_structural_equal(mod, four_matmul_incompatible_batches)


def test_multiple_combine():
    @R.function(private=True)
    def before(
        x1: R.Tensor((2, 1024, 640), "float32"),
        x2: R.Tensor((2, 1024, 640), "float32"),
        w0: R.Tensor((640, 640), "float32"),
        w1: R.Tensor((640, 640), "float32"),
        w2: R.Tensor((640, 640), "float32"),
        w3: R.Tensor((640, 640), "float32"),
        w4: R.Tensor((640, 640), "float32"),
        b0: R.Tensor((640,), "float32"),
        b1: R.Tensor((640,), "float32"),
    ):
        with R.dataflow():
            lv0 = R.matmul(x1, w0)
            lv3 = R.matmul(x2, w3)
            lv1 = R.matmul(x1, w1)
            lv5 = R.add(lv3, b0)
            lv2 = R.matmul(x1, w2)
            lv4 = R.matmul(x2, w4)
            lv6 = R.add(lv4, b1)
            out = (lv0, lv1, lv2, lv5, lv6)
            R.output(out)
        return out

    after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]

    @R.function(private=True)
    def expected(
        x1: R.Tensor((2, 1024, 640), dtype="float32"),
        x2: R.Tensor((2, 1024, 640), dtype="float32"),
        w0: R.Tensor((640, 640), dtype="float32"),
        w1: R.Tensor((640, 640), dtype="float32"),
        w2: R.Tensor((640, 640), dtype="float32"),
        w3: R.Tensor((640, 640), dtype="float32"),
        w4: R.Tensor((640, 640), dtype="float32"),
        b0: R.Tensor((640,), dtype="float32"),
        b1: R.Tensor((640,), dtype="float32"),
    ):
        with R.dataflow():
            lv = R.concat((w0, w1, w2), axis=1)
            lv1 = R.matmul(x1, lv, out_dtype="float32")
            lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=2)
            lv0 = lv2[0]
            lv1_1 = lv2[1]
            lv_1 = R.concat((w3, w4), axis=1)
            lv1_2 = R.matmul(x2, lv_1, out_dtype="float32")
            lv2_1 = R.concat((b0, b1), axis=0)
            lv3 = R.add(lv1_2, lv2_1)
            lv4 = R.split(lv3, indices_or_sections=[640], axis=2)
            lv5 = lv4[0]
            lv2_2 = lv2[2]
            lv6 = lv4[1]
            out = lv0, lv1_1, lv2_2, lv5, lv6
            R.output(out)
        return out

    tvm.ir.assert_structural_equal(after, expected)


def test_check():
    @R.function(private=True)
    def before(
        x1: R.Tensor((2, 1024, 640), "float32"),
        x2: R.Tensor((2, 1024, 640), "float32"),
        w0: R.Tensor((640, 640), "float32"),
        w1: R.Tensor((640, 640), "float32"),
        w2: R.Tensor((640, 640), "float32"),
        w3: R.Tensor((640, 640), "float32"),
        w4: R.Tensor((640, 640), "float32"),
    ):
        with R.dataflow():
            lv0 = R.matmul(x1, w0)
            lv1 = R.matmul(x1, w1)
            lv2 = R.matmul(x1, w2)
            lv3 = R.matmul(x2, w3)
            lv4 = R.matmul(x2, w4)
            out = (lv0, lv1, lv2, lv3, lv4)
            R.output(out)
        return out

    check = lambda *inp: len(inp[1]) > 2  # Ignore branches with two matmuls
    after = CombineParallelMatmul(check)(tvm.IRModule.from_expr(before))["main"]

    @R.function(private=True)
    def expected(
        x1: R.Tensor((2, 1024, 640), dtype="float32"),
        x2: R.Tensor((2, 1024, 640), dtype="float32"),
        w0: R.Tensor((640, 640), dtype="float32"),
        w1: R.Tensor((640, 640), dtype="float32"),
        w2: R.Tensor((640, 640), dtype="float32"),
        w3: R.Tensor((640, 640), dtype="float32"),
        w4: R.Tensor((640, 640), dtype="float32"),
    ):
        with R.dataflow():
            lv = R.concat((w0, w1, w2), axis=1)
            lv1 = R.matmul(x1, lv, out_dtype="float32")
            lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=2)
            lv0 = lv2[0]
            lv1_1 = lv2[1]
            lv2_1 = lv2[2]
            lv3 = R.matmul(x2, w3, out_dtype=None)
            lv4 = R.matmul(x2, w4, out_dtype=None)
            out = (lv0, lv1_1, lv2_1, lv3, lv4)
            R.output(out)
        return out

    tvm.ir.assert_structural_equal(after, expected)


def test_combine_matmul_of_static_and_dynamic_shapes():
    """Combine two matmuls, one with dynamic shape

    The `R.split` operator must have a static list of integer indices
    at which to split the matmul output, because these integer indices
    are stored as operator attributes.  However, the last output can
    still have a dynamic shape.

    """

    @R.function(private=True)
    def before(
        x: R.Tensor((2, 1024, 640), "float32"),
        w0: R.Tensor((640, 640), "float32"),
        w1: R.Tensor((640, "M"), "float32"),
    ):
        M = T.int64()
        with R.dataflow():
            lv0 = R.matmul(x, w0)
            lv1 = R.matmul(x, w1)
            out = (lv0, lv1)
            R.output(out)
        return out

    @R.function(private=True)
    def expected(
        x: R.Tensor((2, 1024, 640), dtype="float32"),
        w0: R.Tensor((640, 640), dtype="float32"),
        w1: R.Tensor((640, "M"), dtype="float32"),
    ) -> R.Tuple(
        R.Tensor((2, 1024, 640), dtype="float32"), R.Tensor((2, 1024, "M"), dtype="float32")
    ):
        M = T.int64()
        with R.dataflow():
            lv: R.Tensor((640, 640 + M), dtype="float32") = R.concat((w0, w1), axis=1)
            lv1: R.Tensor((2, 1024, 640 + M), dtype="float32") = R.matmul(
                x, lv, out_dtype="float32"
            )
            lv2: R.Tuple(
                R.Tensor((2, 1024, 640), dtype="float32"),
                R.Tensor((2, 1024, M), dtype="float32"),
            ) = R.split(lv1, indices_or_sections=[640], axis=2)
            lv0: R.Tensor((2, 1024, 640), dtype="float32") = lv2[0]
            lv1_1: R.Tensor((2, 1024, M), dtype="float32") = lv2[1]
            out: R.Tuple(
                R.Tensor((2, 1024, 640), dtype="float32"),
                R.Tensor((2, 1024, M), dtype="float32"),
            ) = (lv0, lv1_1)
            R.output(out)
        return out

    after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]

    tvm.ir.assert_structural_equal(after, expected)


def test_combine_matmul_of_dynamic_and_static_shapes():
    """Combine two matmuls, one with dynamic shape

    Like `test_combine_matmul_of_static_and_dynamic_shapes`, but the
    dynamic-shaped matmul is encountered first.  Due to the
    requirements imposed by `R.split` storing the split indices as
    static integers, the static-shaped weights must occur first in the
    concatenated weights.
    """

    @R.function(private=True)
    def before(
        x: R.Tensor((2, 1024, 640), "float32"),
        w0: R.Tensor((640, "M"), "float32"),
        w1: R.Tensor((640, 640), "float32"),
    ):
        M = T.int64()
        with R.dataflow():
            lv0 = R.matmul(x, w0)
            lv1 = R.matmul(x, w1)
            out = (lv0, lv1)
            R.output(out)
        return out

    @R.function(private=True)
    def expected(
        x: R.Tensor((2, 1024, 640), dtype="float32"),
        w0: R.Tensor((640, "M"), dtype="float32"),
        w1: R.Tensor((640, 640), dtype="float32"),
    ) -> R.Tuple(
        R.Tensor((2, 1024, "M"), dtype="float32"), R.Tensor((2, 1024, 640), dtype="float32")
    ):
        M = T.int64()
        with R.dataflow():
            lv: R.Tensor((640, 640 + M), dtype="float32") = R.concat((w1, w0), axis=1)
            lv1: R.Tensor((2, 1024, 640 + M), dtype="float32") = R.matmul(
                x, lv, out_dtype="float32"
            )
            lv2: R.Tuple(
                R.Tensor((2, 1024, 640), dtype="float32"),
                R.Tensor((2, 1024, M), dtype="float32"),
            ) = R.split(lv1, indices_or_sections=[640], axis=2)
            lv0: R.Tensor((2, 1024, M), dtype="float32") = lv2[1]
            lv1_1: R.Tensor((2, 1024, 640), dtype="float32") = lv2[0]
            out: R.Tuple(
                R.Tensor((2, 1024, M), dtype="float32"),
                R.Tensor((2, 1024, 640), dtype="float32"),
            ) = (lv0, lv1_1)
            R.output(out)
        return out

    after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]

    tvm.ir.assert_structural_equal(after, expected)


def test_limit_one_dynamic_shape_in_combined_matmul():
    """Combine two matmuls, one with dynamic shape

    Like `test_combine_matmul_of_static_and_dynamic_shapes`, but with
    two dynamic weights that could, in principle, be merged together.
    Because `R.split` must have integer indices at which to split,
    only one of the dynamic outputs can be part of the combined
    matmul.
    """

    @R.function(private=True)
    def before(
        x: R.Tensor((2, 1024, 640), "float32"),
        w0: R.Tensor((640, "M"), "float32"),
        w1: R.Tensor((640, 640), "float32"),
        w2: R.Tensor((640, "N"), "float32"),
    ):
        M = T.int64()
        with R.dataflow():
            lv0 = R.matmul(x, w0)
            lv1 = R.matmul(x, w1)
            lv2 = R.matmul(x, w2)
            out = (lv0, lv1, lv2)
            R.output(out)
        return out

    @R.function(private=True)
    def expected(
        x: R.Tensor((2, 1024, 640), dtype="float32"),
        w0: R.Tensor((640, "M"), dtype="float32"),
        w1: R.Tensor((640, 640), dtype="float32"),
        w2: R.Tensor((640, "N"), "float32"),
    ) -> R.Tuple(
        R.Tensor((2, 1024, "M"), dtype="float32"),
        R.Tensor((2, 1024, 640), dtype="float32"),
        R.Tensor((2, 1024, "N"), dtype="float32"),
    ):
        M = T.int64()
        with R.dataflow():
            concat_weights = R.concat((w1, w0), axis=1)
            concat_output = R.matmul(x, concat_weights, out_dtype="float32")
            split_output: R.Tuple(
                [R.Tensor([2, 1024, 640], dtype="float32"), R.Tensor([2, 1024, M], dtype="float32")]
            ) = R.split(concat_output, indices_or_sections=[640], axis=2)
            lv0 = split_output[1]
            lv1 = split_output[0]
            lv2 = R.matmul(x, w2)
            out = (lv0, lv1, lv2)
            R.output(out)
        return out

    after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]

    tvm.ir.assert_structural_equal(after, expected)


if __name__ == "__main__":
    tvm.testing.main()