masahi commented on code in PR #13050:
URL: https://github.com/apache/tvm/pull/13050#discussion_r995451386
##########
src/relay/backend/task_extraction.cc:
##########
@@ -42,26 +46,35 @@ Array<meta_schedule::ExtractedTask> ExtractTask(IRModule
mod, Target target,
mod = transform::Sequential(pass_seqs)(std::move(mod));
std::vector<ExtractedTask> tasks;
- std::unordered_map<tec::CCacheKey, ExtractedTask> cache;
+
+ auto mod_eq = meta_schedule::ModuleEquality::Create(mod_eq_name);
+
+ std::unordered_map<IRModule, ExtractedTask, ModuleHash, ModuleEqual> cache(
+ /*bucket_count*/ 0, ModuleHash(*mod_eq), ModuleEqual(*mod_eq));
+
PostOrderVisit(mod->Lookup("main"), [&target, &tasks, &cache,
&tir_converter](const Expr& exp) {
if (exp->IsInstance<FunctionNode>()) {
Function relay_func = Downcast<Function>(exp);
if (!relay_func->HasNonzeroAttr(attr::kPrimitive)) {
return;
}
- tec::CCacheKey cache_key(relay_func, target);
- auto it = cache.find(cache_key);
- if (it != cache.end()) {
- it->second->weight += 1;
- return;
- }
+
auto [inputs_outputs, constants, fused_name] =
tec::LowerTECompute(relay_func, target, /*return_inputs=*/true);
+
if (Optional<tir::PrimFunc> f = tir_converter(inputs_outputs,
constants)) {
+ IRModule tir_mod = PrimFuncToIRModule(f.value());
+
+ auto it = cache.find(tir_mod);
+ if (it != cache.end()) {
+ it->second->weight += 1;
+ return;
+ }
+
IRModule relay_mod({{GlobalVar(fused_name), relay_func}});
- ExtractedTask task(fused_name, relay_mod, target,
{PrimFuncToIRModule(f.value())}, 1);
+ ExtractedTask task(fused_name, relay_mod, target, {tir_mod}, 1);
tasks.push_back(task);
- cache.emplace(cache_key, task);
+ cache.emplace(tir_mod, task);
Review Comment:
This is a great point. In my branch where I'm working on anchor block
tuning, I do need to be careful here. When there are both `conv2d` and `conv2d
+ add`, we need to extract `conv2d`. To do that, before I put modules into the
cache, I sort them based on the fused function name, so that `conv2d` comes
before `conv2d + add`.
This way, after `conv2d` is put to the cache, we can automatically skip
adding `conv2d + add`.
--
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]
