comaniac opened a new pull request #8560: URL: https://github.com/apache/tvm/pull/8560
Per discussion in https://discuss.tvm.apache.org/t/te-compiler-device-type-error-in-updatemainworkspacesize/10536/5, the updated TE compiler has a different behavior as the old compile engine in terms of visiting each Relay function for compilation. Specifically, TE compiler now makes sure the function name is unique when it appears multiple times in a model. This breaks the current task extraction mechanism, which uses function name and TE compute hash as the key to de-duplicate tasks. As a result, the same task will be extracted multiple times, and all tasks have the weight 1. This PR makes the following changes to resolve this issue: 1. Change the type of `wkl_key_to_weight` from `Dict[str, int]` to `Dict[str, Tuple[int, Set[Str]]]`, mapping from a workload key to the weight as well as associated function names. In this way, we guarantee that functions with the same TE compute will be extracted once. 2. Add `func_name_to_wkl_key`, mapping from the unique function name to its workload key. When TE compiler is used and `te_compiler_update_weights` is invoked, it obtains the workload key of the given function and aggregate the weights. 3. Refactor unit tests to be pytest parameterized. Also improve the control flow test to have 2 identical Conv2Ds. Without this PR, this test will fail because it falls back to use VM compiler due to the control flow and extracts 2 identical tasks. cc @masahi @csullivan @jroesch @merrymercy -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
