[TVMC] Apply constant folding when converting layout

[lhutton1]

This commit ensures that constant folding is applied when a desired layout is selected during compilation. It ensures that layout_transform operations are removed where possible so that pattern matching for BYOC backends can work effectively. A test has been added.

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @Mousius, @gromero _{See #10317 for details}
• Built docs for commit 792ff1a can be found here.

_{Generated by tvm-bot}

[lhutton1]

also cc @leandron @ekalda

[Mousius]

This seems to be re-testing the existing passes? Instead we could just test that it's using them? I don't know if there's any off the shelf infrastructure to get a list of passes but we could do something like:

mock_sequence = MagicMock()
mock_pass_runner = MagicMock()
mock_sequence.return_value = mock_pass_runner
monkeypatch.setattr(transform, 'Sequential', mock_sequence);

convert_graph_layout(mod, "NHWC")

# Might need to mock out a bunch of passes here with some
mock_sequence.assert_called_once_with([
 relay.transform.RemoveUnusedFunctions(),
 relay.transform.ConvertLayout({
        "nn.conv2d": ["NHWC", "default"],
        "nn.conv2d_transpose": ["NHWC", "default"],
        "qnn.conv2d": ["NHWC", "default"],
 }),
 relay.transform.FoldConstant(),
])
mock_pass_runner.assert_called_once_with(mod);

[ekalda]

The test in the patch makes sense to me, it checks that the constant has been folded which wouldn't happen unless we have run the FoldConstant pass. I suppose mocking would work as well, but don't we then have the problem of having to keep the dict of ops and preferred layouts in the codebase manually in sync with the same dict in the test?

[lhutton1]

I think the issue with the current test is that it checks more than it really needs to and might become a maintenance burden going forwards. Ideally I'd like to make use of pass instruments to check that ConvertLayout and FoldConstant are run in the correct order, which means that we aren't tied to the implementation details of the passes. However, the use of PassContext here overwrites any context set by the user (and therefore by the test that adds the instrumentation) so I think we should remove this and set the context from the callee(s) - given the test works for now, could we do this in a follow up?

[ekalda]

Yes, I agree, that testing the implementation details of the passes themselves in not ideal. I'd be fine with merging the current patch as it is and improving the test in a follow up, once we have cleared up the nested PassContexts if @Mousius agrees?

[Mousius]

Agree, this will fail if we disable constant folding so that'll work for now 😸

[ekalda]

LGTM! :)

[ekalda]

The test in the patch makes sense to me, it checks that the constant has been folded which wouldn't happen unless we have run the FoldConstant pass. I suppose mocking would work as well, but don't we then have the problem of having to keep the dict of ops and preferred layouts in the codebase manually in sync with the same dict in the test?