refactor

pfk-beta · Mar 11, 2022 · 74636be · 74636be
1 parent 99f1701
commit 74636be
Showing 1 changed file with 10 additions and 15 deletions.
diff --git a/src/meta_schedule/integration.cc b/src/meta_schedule/integration.cc
@@ -26,8 +26,8 @@ namespace tvm {
 namespace meta_schedule {
 
 /**************** Utility functions ****************/
-template <class FunctionType>
-Optional<GlobalVar> GetOnlyOneFunctionKey(const IRModule& mod) {
+template <class FunctionType, class RetType, class Callback>
+Optional<RetType> GetOnlyOneFunctionCommon(const IRModule& mod, Callback on_found) {
   if (mod->functions.size() != 1) {
     return NullOpt;
   }
@@ -36,26 +36,21 @@ Optional<GlobalVar> GetOnlyOneFunctionKey(const IRModule& mod) {
     if (!func->IsInstance<typename FunctionType::ContainerType>()) {
       return NullOpt;
     } else {
-      return kv.first;
+      return on_found(kv);
     }
   }
   return NullOpt;
 }
 
+template <class FunctionType>
+Optional<GlobalVar> GetOnlyOneFunctionKey(const IRModule& mod) {
+  return GetOnlyOneFunctionCommon<FunctionType, GlobalVar>(mod, [](auto kv) { return kv.first; });
+}
+
 template <class FunctionType>
 Optional<FunctionType> GetOnlyOneFunction(const IRModule& mod) {
-  if (mod->functions.size() != 1) {
-    return NullOpt;
-  }
-  for (const auto& kv : mod->functions) {
-    const BaseFunc& func = kv.second;
-    if (!func->IsInstance<typename FunctionType::ContainerType>()) {
-      return NullOpt;
-    } else {
-      return Downcast<FunctionType>(func);
-    }
-  }
-  return NullOpt;
+  return GetOnlyOneFunctionCommon<FunctionType, FunctionType>(
+      mod, [](auto kv) { return Downcast<FunctionType>(kv.second); });
 }
 
 template <class FunctionType>