[MetaSchedule] Allow skipping exact NDArray rewrite in RemoveWeightLayoutRewriteBlock

[masahi]

I found that, using MS evo search + RewriteLayout + link-params=True, where the latter combination is supported by #12991 and #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

tvm/src/meta_schedule/search_strategy/evolutionary_search.cc

Lines 206 to 209 in 6780c9f

	for (const Schedule& sch : picks) {
	measure_inputs.push_back(
	MeasureCandidate(sch, ArgInfo::FromEntryFunc(sch->mod(), /*remove_preproc=*/true)));
	}

-> ArgInfo::FromEntryFun

tvm/src/meta_schedule/arg_info.cc

Lines 104 to 107 in 111169c

	Array<ArgInfo> ArgInfo::FromEntryFunc(const IRModule& mod, bool remove_preproc) {
	if (remove_preproc) {
	IRModule new_mod = tir::transform::RemoveWeightLayoutRewriteBlock()(mod);
	return ArgInfo::FromPrimFunc(FindEntryFunc(new_mod));

-> RemoveWeightLayoutRewriteBlock

tvm/src/tir/transforms/remove_weight_layout_rewrite_block.cc

Lines 163 to 167 in 189338c

	Stmt VisitStmt_(const AllocateConstNode* alloc) final {
	if (auto it = buffer_var_to_index_map_.find(alloc->buffer_var.get());
	it != buffer_var_to_index_map_.end()) {
	auto new_body = StmtMutator::VisitStmt(alloc->body);
	auto rewritten_ndarray = it->second->MapNDArray(alloc->data.value());

-> 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.

[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 @Hzfengsy, @junrushao _{See #10317 for details}
• Built docs for commit 02da79e can be found here.

_{Generated by tvm-bot}