Function test_nested_diamond

tests/python/relax/test_dataflow_pattern.py:725–805

()

def test_nested_diamond():
    @tvm.script.ir_module
    class DiamondInDiamond:
        @R.function
        def main(x: R.Tensor((32, 32), "float32"), w: R.Tensor((32, 32), "float32")) -> R.Tensor:
            with R.dataflow():
                #   matmul0      matmul1
                #     /    \    /    \
                # sigmoid2  add4  sigmoid3
                #     \    /    \    /
                #      add5      add6
                #          \    /
                #           add7
                lv0 = R.call_dps_packed(
                    "extern_matmul", (x, w), R.Tensor((32, 32), dtype="float32")
                )
                lv1 = R.call_dps_packed(
                    "extern_matmul", (x, w), R.Tensor((32, 32), dtype="float32")
                )
                lv2 = R.call_dps_packed(
                    "extern_sigmoid", (lv0), R.Tensor((32, 32), dtype="float32")
                )
                lv3 = R.call_dps_packed(
                    "extern_sigmoid", (lv1), R.Tensor((32, 32), dtype="float32")
                )
                lv4 = R.call_dps_packed(
                    "extern_add", (lv0, lv1), R.Tensor((32, 32), dtype="float32")
                )
                lv5 = R.call_dps_packed(
                    "extern_add", (lv2, lv4), R.Tensor((32, 32), dtype="float32")
                )
                lv6 = R.call_dps_packed(
                    "extern_add", (lv3, lv4), R.Tensor((32, 32), dtype="float32")
                )
                lv7 = R.call_dps_packed(
                    "extern_add", (lv5, lv6), R.Tensor((32, 32), dtype="float32")
                )
                R.output(lv7)
            return lv7

    # match matmul0 diamond
    with PatternContext() as ctx:
        sigmoid2 = is_call_dps_packed("extern_sigmoid")
        add4 = is_call_dps_packed("extern_add")
        is_call_dps_packed("extern_matmul").fork_to(sigmoid2, add4)
        add5 = is_call_dps_packed("extern_add")
        sigmoid2 >> add5
        add4 ^ add5
        assert ctx.match_dfb(DiamondInDiamond["main"].body.blocks[0])

    # counter case: mis-match matmul0 diamond
    with PatternContext() as ctx:
        sigmoid2 = is_call_dps_packed("extern_sigmoid")
        add4 = is_call_dps_packed("extern_add")
        is_call_dps_packed("extern_matmul").fork_to(sigmoid2, add4)
        add5 = is_call_dps_packed("extern_add")
        sigmoid2 >> add5
        add4 >> add5  # not only-used-by relation