[TVMC][UMA] Support using UMA with TVMC

[rafzi]

This change introduces:

• A TVMExtension interface to extend TVMC with custom code. Currently its only use is to add UMABackends.
• A command line option --experimental-tvm-extension <dir> to point to extension modules.
• An envvar $TVM_EXTENSION_DIR for the same purpose.

See discussion thread: https://discuss.tvm.apache.org/t/tvmc-uma-integration/14309

@Mousius @cgerum @MichaelJKlaiber @PhilippvK

[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, @leandron, @lhutton1 _{See #10317 for details}

_{Generated by tvm-bot}

[lhutton1]

This looks really nice, thanks @rafzi :) I had a couple of questions and some suggestions for tests

[lhutton1]

Curious, why do we need to remove help option?

[rafzi]

I encountered an issue where the TVMCSuppressedArgumentParser still caused a SystemExit exception. However, I cannot seem to reproduce the issue right now. So I'll remove this again.

[rafzi]

Alright, I found the input again: "tvmc compile --help". When removing the help option from the subparser, this is resolved. The help option still works because it is enabled through the main parser.

[lhutton1]

I tried this out locally but see some difference in the help prompt when executing tvmc compile --help. Before the patch we get something like:

$ tvmc compile --help
usage: tvmc compile [-h] [--cross-compiler CROSS_COMPILER]
                    [--cross-compiler-options CROSS_COMPILER_OPTIONS]
                    [--desired-layout {NCHW,NHWC}] [--dump-code FORMAT]
                    [--model-format {keras,onnx,pb,tflite,pytorch,paddle,relay}]
                    [-o OUTPUT] [-f {so,mlf}] [--pass-config name=value]
                    [--target TARGET]
...
positional arguments:
  FILE                  path to the input model file.

optional arguments:
  -h, --help            show this help message and exit
  --cross-compiler CROSS_COMPILER
                        the cross compiler to generate target libraries, e.g.
                        'aarch64-linux-gnu-gcc'.
  --cross-compiler-options CROSS_COMPILER_OPTIONS
                        the cross compiler options to generate target
                        libraries, e.g. '-mfpu=neon-vfpv4'.
  --desired-layout {NCHW,NHWC}
                        change the data layout of the whole graph.

While after seems to lead to:

$ tvmc compile --help
usage: tvmc compile [--experimental-tvm-extension EXPERIMENTAL_TVM_EXTENSION]
                    [--cross-compiler CROSS_COMPILER]
                    [--cross-compiler-options CROSS_COMPILER_OPTIONS]
                    [--desired-layout {NCHW,NHWC}] [--dump-code FORMAT]
                    [--model-format {keras,onnx,pb,tflite,pytorch,paddle,relay}]
                    [-o OUTPUT] [-f {so,mlf}] [--pass-config name=value]
                    [--target TARGET]
...
tvmc compile: error: the following arguments are required: FILE

Could you confirm if you see the same behaviour? I think it would be really helpful if we can keep the argument descriptions as before

[rafzi]

Alright, I can confirm this issue. The underlying issue is that the TVMCSuppressedArgumentParser still throws SystemExit when given such a command line. We also cannot override error of ArgumentParser, because the failing parser is the subparser, which is not a TVMCSuppressedArgumentParser.

Before this patch, this problem already occurs. For example with tvmc run --help, SystemExit is thrown and a partial help message is printed (missing --project-option). So my suggested fix is to move the parse_known_args as far down as possible to yield a mostly complete help message.

[lhutton1]

Thanks for investigating! This looks much better, so I won't block on it. Perhaps we can create a separate issue to track this bug?

[lhutton1]

would be good to add a test for this case

[rafzi]

I tried adding this, but encountered an issue with race conditions. The way the sys.path is currently modified by extensions.py is problematic when more than one test is run in parallel. Is there already an approach for serializing tests in TVM? There doesn't seem to be an upstream solution: pytest-dev/pytest-xdist#18

[lhutton1]

Could we instead avoid modifying sys.path when loading a module? I had a naive search and wondered if this could help? cc @Mousius @leandron who might know more about these things

[rafzi]

While such an approach should work for single source files, I'm not sure how that would work for packages or namespace packages.

[lhutton1]

Apologies for the delay, I had a think about this but wasn't able to come up with a better solution myself, so I won't block on this

[cgerum]

Thanks @rafzi for the proposal, thanks @lhutton1 for the feedback.

I also have a few minor requests.

[MichaelJKlaiber]

I agreed with @lhutton1 and @cgerum 's comments.
Besides that everything LGTM.
Thanks @rafzi for the contribution

[PhilippvK]

@tvm-bot rerun

[rafzi]

Sorry for the late update.

The outstanding test failure is likely because of the outstanding issue that the sys.path is modified in order to support all module loading. This probably inteferres with running the tests in parallel.

https://github.com/tum-ei-eda/tvm/blob/c1c2d3ca2db1c56d3cc8d695530ab7aa8b90ea00/python/tvm/driver/tvmc/extensions.py#L24-L34

I'm not sure how to resolve this as I did not find another method to achieve the same kind of module loading. Do we need to compromise here and settle for a simple module load? I unresolved the relevant discussion with @lhutton1.

@cgerum @PhilippvK