BYOC Partitioning Support on TVMC

[leandron]

BYOC Partitioning Support on TVMC

In this RFC, we are proposing an implementation that allow users to reach BYOC targets using TVMC. This is done by applying the expected sequence of passes, in the order requested by the user via command line.

Current Status

In the current implementation, users can’t target BYOC backends directly using TVM Target API, needing to resort to custom passes execution using Python custom scripts.

Proposed solution

This RFC proposes a mechanism to expose BYOC targets to end-users via TVMC, as a way to offer that feature via command line. It does so by linking the target definitions passed by --target, with the series of passes and partitioning that currently needs to be done manually.

In the specific case of BYOC, TVMC needs to invoke manual partitioning and a set of passes that will adjust the graph for our expected BYOC.

One or more BYOC implementations are expected to be targeted - together or separate - in any (reasonable) combination the user might want.

To accomplish this, the present RFC expects to extend the existing --target option syntax, to accommodate the target(s), in a way to support BYOC, as well as keeping the target option as a consistent way to define a target.

Example:

tvmc compile --target="ethos-n77, llvm -mtriple=aarch64-linux-gnu -mattr=+neon" ...
                       ^^^^^^^^^

The sets of passes to be executed for a given BYOC are implemented by a class, that will link the command line option with that BYOC target. This class also will be able to translate the options given via command line, to be translated to a dictionary passed to the tvm.transform.PassContext, at compilation time.

Next steps

Together with this RFC, I’ll soon submit a PR https://github.com/apache/tvm/pull/7304 that implements what is here. We are currently using this approach successfully to test the codegen for Arm Ethos-N NPU. This will also be included the PR.

@mjs @comaniac @tqchen @ramana-arm @mbaret

[comaniac]

Thanks for the RFC. I have the following questions in the example.

1. Looks like the separator is comma, but how would you differentiate the target options, such as -mattr=+neon,fp-armv8,thumb-mode ?

2. How would you recognize which target in this string is the TVM backend (e.g., llvm, cuda)?

In general, we should first improve the composite target to naturally support BYOC backends instead of having the target dispatching logic in a command line module, but unfortunately we haven’t really got a chance to work on it in the past months.

Also cc @zhiics

[leandron]

Hi @comaniac, sorry it took me a while to get back here.

@comaniac

Looks like the separator is comma, but how would you differentiate the target options, such as -mattr=+neon,fp-armv8,thumb-mode ?

So, in this case, the parser would be happy with quotes to represent the substring. In an extended version of your example, it would be ethos-n77, llvm -mattr='+neon,fp-armv8,thumb-mode'.

@comaniac

How would you recognize which target in this string is the TVM backend (e.g., llvm , cuda )?

There are some validations in place (and we can extend it to validate more cases, once identified).

The order in which you declare your targets is respected, so naturally, the TVM target that is the fallback for everything is the last one - this is one of the validations in place.

[comaniac]

Thanks for the clarification.

So, in this case, the parser would be happy with quotes to represent the substring. In an extended version of your example, it would be ethos-n77, llvm -mattr='+neon,fp-armv8,thumb-mode' .

How about just use the other separator like ; so that it would be clearer?

There are some validations in place (and we can extend it to validate more cases, once identified). The order in which you declare your targets is respected, so naturally, the TVM target that is the fallback for everything is the last one - this is one of the validations in place.

Understood, although I intuitively attempt to put the TVM backend in the beginning instead of in the end lol. In the short term, I guess it’s fine to just explicitly mention the last one is the fallback target in the help message.

[mjs]

The primary goal of tvmc is to provide a command line interface to tvm. That means its primary use will be via shell, either interactive or in script. Therefore giving special meaning in the tvmc command line to any character that the shell gives special meaning to creates an awkward usage forcing the user to inject addition quoting or escaping. “;” is the command separator in shell, which makes it less than ideal to have meaning within a tvmc option value. Instead better to focus on characters that don’t have special shell meaning especially in common use cases.

[manupa-arm]

Hi @comaniac,

Understood, although I intuitively attempt to put the TVM backend in the beginning instead of in the end lol. In the short term, I guess it’s fine to just explicitly mention the last one is the fallback target in the help message.

I think we were thinking of ordering that could be usable when we can support multiple BYOC codegens. Say for e.g., when we want to do Ethos-N77 then Arm Compute Library and fallback being for llvm, I think its intuitive to have {ethos-n77, acl, llvm} compared to the alternatives of {llvm, ethos-n77, acl} or {llvm, acl, ethos-n77} – I mean descending priority make sense to me vs ascending/mixed priorities. What do you think ?

I agree with @mjs we should avoid escaping as much as possible in the shell.

[comaniac]

Fair enough. Let’s follow the order you proposed.

@mjs is right about the escaping “;” that I missed before. I think it should be fine for now to just use the comma. One the composite target support is landed, all these issues would be gone, although I currently don’t have a timeline for it.

[leandron]

@comaniac

Fair enough. Let’s follow the order you proposed.

I updated the patch, that covers a few use cases that people might just expect it to work without nested quotes. Commas are still used as delimiter.

The examples below are expected to work with the latest version:

# no nested quotes required is everything is together
--target="ethos-n77, llvm -mattr=+neon,fp-armv8,thumb-mode"

# nested quotes required if you have spaces though
--target="ethos-n77, llvm -mattr='+neon, fp-armv8, thumb-mode'"

I added a few more test cases to cover other corner cases as well.

[leandron]

This is now merged in [TVMC] Add composite target passes for compilation and tuning (#7304) · apache/tvm@d16f282 · GitHub

Thanks

[jtuyls]

@leandron I had a I look at adding Vitis AI as a target to TVMC and we will need to make a couple of changes for that in our codegen/runtime to use the ‘config_key’ in the way you are doing it for Ethos-N. Also, we will need a way of passing config information to the partition_function because the partitioning might depend on the the specific target Vitis AI DPU, which is passed as a config. I think we can pass the target options (**codegen_from_cli["opts"]) to the partition_function here: tvm/compiler.py at a1118039ea80e4a4e79fb28012b375867b0a8613 · apache/tvm · GitHub so that the partitioning function can potentially make use of those. I will be working on a patch to add Vitis AI so let me know if you have any thoughts on this.

[leandron]

That’s great news @jtuyls!

@jtuyls

we will need to make a couple of changes for that in our codegen/runtime to use the ‘config_key’ in the way you are doing it for Ethos-N.

I think it would be great to have as many options as possible being passed via the PassContext, as it makes it simpler to have generic TVM support, but agree that in some cases, we might need to send parameters to the partition function itself.

@jtuyls

I think we can pass the target options (**codegen_from_cli["opts"]) to the partition_function here: tvm/compiler.py at a1118039ea80e4a4e79fb28012b375867b0a8613 · apache/tvm · GitHub so that the partitioning function can potentially make use of those.

This codegen_from_cli["opts"] will contain everything you send as options to your codegen, such as ethos-n77 -myopt=value -option2=value2, you get { "myopt": "value", "option2": "value2" }.

In this opts dictionary we may have things that are aimed at the PassContext and, as you describe, things that are aimed at the partition function, so I think we will need to rework the partition calls with mixed arguments, to follow a single “kwargs” signature.

Having a **opts might look opaque, but I think it is safer for the function to check with config_i_need = opts.get('key', default). For the sake a documentation, I think we should add into the docstring the options we expect to use in the partition function, as well as the defaults.

So, as a practical example, taking the current Compute Library partition function, it would become:

def partition_for_arm_compute_lib(mod, params=None, **opts):
                                                    ^^^^^^

I think it would be great to have a standard signature for the partition functions that look like this above.

Happy to hear other suggestions and review your PR once you have something.

PS: there is also a similar call to be adjusted in tvm/autotuner.py at a1118039ea80e4a4e79fb28012b375867b0a8613 · apache/tvm · GitHub