dedup tasks

apache · Mar 11, 2022 · dfaf496 · dfaf496
1 parent e49d500
commit dfaf496
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/src/relay/backend/task_extraction.cc b/src/relay/backend/task_extraction.cc
@@ -46,10 +46,12 @@ Array<ExtractedTask> ExtractTask(IRModule mod, Target target, Map<String, Consta
   auto opt_mod = seq(std::move(mod));
 
   Array<ExtractedTask> tasks;
-  PostOrderVisit(opt_mod->Lookup("main"), [target, &tasks](const Expr& exp) {
+  std::unordered_set<tec::CCacheKey> cache_;
+  PostOrderVisit(opt_mod->Lookup("main"), [target, &tasks, &cache_](const Expr& exp) {
     if (exp->IsInstance<FunctionNode>()) {
       Function relay_func = Downcast<Function>(exp);
-      if (relay_func->HasNonzeroAttr(attr::kPrimitive)) {
+      tec::CCacheKey cache_key(relay_func, target);
+      if (relay_func->HasNonzeroAttr(attr::kPrimitive) && cache_.find(cache_key) == cache_.end()) {
         Array<te::Tensor> outputs;
         std::string fused_name;
         std::tie(outputs, fused_name) =
@@ -59,6 +61,7 @@ Array<ExtractedTask> ExtractTask(IRModule mod, Target target, Map<String, Consta
         auto relay_mod = IRModule({{prim_fn_var, relay_func}});
         auto tir_mod = IRModule({{prim_fn_var, prim_func}});
         tasks.push_back(ExtractedTask(prim_fn_var->name_hint, relay_mod, target, {tir_mod}));
+        cache_.insert(cache_key);
       }
     }
   });

diff --git a/src/relay/backend/te_compiler_cache.cc b/src/relay/backend/te_compiler_cache.cc
@@ -757,7 +757,7 @@ CachedFunc ShapeFuncFor(const Function& prim_func, const Target& target,
 std::pair<Array<te::Tensor>, std::string> LowerTECompute(const Function& source_func, Target target,
                                                          bool return_inputs) {
   LowerToTECompute lower_te_compute(target);
-  auto outputs = lower_te_compute.Lower(source_func, [&](std::string name) { return name; });
+  auto outputs = lower_te_compute.Lower(source_func, [&](std::string name) { return name;});
   // Following ScheduleBuilder, remove placeholder ops from outputs.
   tvm::Array<te::Tensor> tensor_outs;
   for (const auto& tensor : outputs) {