From 8108f596008289775bad2894eb06e3e0b891a3e9 Mon Sep 17 00:00:00 2001
From: Cody Yu <comaniac0422@gmail.com>
Date: Tue, 23 Jun 2020 08:23:58 -0700
Subject: [PATCH] [Relay]Allow every runtime module to handle constants (#5885)

* update source module

* address comment
---
 src/target/source/source_module.cc | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/target/source/source_module.cc b/src/target/source/source_module.cc
index 274711183ab9..68a34c8304f8 100644
--- a/src/target/source/source_module.cc
+++ b/src/target/source/source_module.cc
@@ -58,15 +58,18 @@ runtime::Module CreateMetadataModule(
   // Wrap all submodules in the initialization wrapper.
   std::unordered_map<std::string, std::vector<std::string>> sym_metadata;
   for (runtime::Module it : modules) {
-    CHECK_EQ(it->type_key(), "c") << "Only csource submodule is handled for now";
-    String symbol = it.GetFunction("get_symbol")();
-    Array<String> variables = it.GetFunction("get_const_vars")();
-    std::vector<std::string> arrays;
-    for (size_t i = 0; i < variables.size(); i++) {
-      arrays.push_back(variables[i].operator std::string());
+    auto pf_sym = it.GetFunction("get_symbol");
+    auto pf_var = it.GetFunction("get_const_vars");
+    if (pf_sym != nullptr && pf_var != nullptr) {
+      String symbol = pf_sym();
+      Array<String> variables = pf_var();
+      std::vector<std::string> arrays;
+      for (size_t i = 0; i < variables.size(); i++) {
+        arrays.push_back(variables[i].operator std::string());
+      }
+      CHECK_EQ(sym_metadata.count(symbol), 0U) << "Found duplicated symbol: " << symbol;
+      sym_metadata[symbol] = arrays;
     }
-    CHECK_EQ(sym_metadata.count(symbol), 0U) << "Found duplicated symbol: " << symbol;
-    sym_metadata[symbol] = arrays;
   }
 
   // Wrap the modules.