masahi opened a new pull request, #13052: URL: https://github.com/apache/tvm/pull/13052
I found that, using MS evo search + `RewriteLayout` + `link-params=True`, where the latter combination is supported by https://github.com/apache/tvm/pull/12991 and https://github.com/apache/tvm/pull/12949, is currently broken for the following reasons: * The manual NDArray rewriting via `IndexMap::MapNDArray` is slow * `IndexMap::MapNDArray` is called more than thousands of times during evo search **per iteration**, by the following call sequence: `AssembleCandidates` https://github.com/apache/tvm/blob/6780c9f87db6620409f8f58c2c2925c7bd7b6681/src/meta_schedule/search_strategy/evolutionary_search.cc#L206-L209 -> `ArgInfo::FromEntryFun` https://github.com/apache/tvm/blob/111169c7df2831ab8ee40d5388ebcfcf551fd86f/src/meta_schedule/arg_info.cc#L104-L107 -> `RemoveWeightLayoutRewriteBlock` https://github.com/apache/tvm/blob/189338c919c0876cf5909b99fb125ee2a7fbe2c6/src/tir/transforms/remove_weight_layout_rewrite_block.cc#L163-L167 -> `MapNDArray` So what happens right now is that evo search + `link-params=True` would make tuning appear hanged, while it is busy doing thousands of `MapNDArray` in `AssembleCandidates`. My proposed workaround for this problem is based on the observation that, when `RemoveWeightLayoutRewriteBlock` is called by `FromEntryFun` during evo search, the actual content of NDArray before and after rewrite does not matter. So we don't have to use `MapNDArray` during tuning. To enable skipping such "exact" NDArray rewrite, I'm introducing a flag to `RemoveWeightLayoutRewriteBlock`. The exact rewrite by `MapNDArray` is only required when the pass is called from TECompiler during the final compilation. -- 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]
