Method test

exir/backend/test/test_backends_nested.py:218–258

Partitions the cond ops into the delegate

(self)

class TestNestedBackends(unittest.TestCase):
    def test(self) -> None:
        """
        Partitions the cond ops into the delegate
        """

        m = M()
        orig_res = m(*m.get_example_inputs())
        orig = to_edge(
            export(m, m.get_example_inputs(), strict=True),
            compile_config=exir.EdgeCompileConfig(_check_ir_validity=False),
        )

        partitioned = orig
        partitioned._edge_programs["forward"] = to_backend(
            orig.exported_program(), Backend1PartitionerDemo()
        )

        new_res = partitioned.exported_program().module()(*m.get_example_inputs())[0]
        self.assertTrue(torch.allclose(orig_res, new_res))

        # The toplevel module should have lowered the cond and add op
        toplevel_lowered = get_lowered_submodules(
            partitioned.exported_program().graph_module
        )
        self.assertEqual(len(toplevel_lowered), 1)
        toplevel_lowered = toplevel_lowered[0][1]
        self.maxDiff = None
        self.assertEqual(
            str(toplevel_lowered.processed_bytes),
            (
                'b"Backend1::('
                + "call_delegate(b'Backend2::aten.add.Tensor;');"
                + "cond("
                # True function of toplevel cond (nested cond)
                + "cond(call_delegate(b'Backend2::aten.add.Tensor;aten.mm.default;');,call_delegate(b'Backend2::aten.mm.default;'););"
                # True function of toplevel cond (delegated add)
                + "call_delegate(b'Backend2::aten.add.Tensor;');,"
                # False function of toplevel cond
                + "call_delegate(b'Backend2::aten.mm.default;'););)\""
            ),
        )