Fix a few bugs: (#13)

1. Don't add relay main function to list of lowered TIR functions 2. Don't skip visiting call to relay function in graph runtime codegen
jroesch · Mar 19, 2021 · 755f5e1 · 755f5e1
1 parent 7ce8f66
commit 755f5e1
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 14 deletions.
diff --git a/python/tvm/relay/backend/_backend.py b/python/tvm/relay/backend/_backend.py
@@ -80,7 +80,6 @@ def build(mod, target, target_host=None):
     """
     if target_host == "":
         target_host = None
-    import pdb; pdb.set_trace()
     return tvm.driver.build(mod, target=target, target_host=target_host)
 
 

diff --git a/src/relay/backend/graph_runtime_codegen.cc b/src/relay/backend/graph_runtime_codegen.cc
@@ -254,8 +254,6 @@ class GraphRuntimeCodegen : public backend::MemoizedExprTranslator<std::vector<G
 
     ret.lowered_funcs = lowered_module.per_target_module;
     std::cout << "Modules: " << ret.lowered_funcs << std::endl;
-    auto it = ret.lowered_funcs.begin();
-    (*it).second->Update(main_module);
     ret.external_mods = lowered_module.external_mods;
     return ret;
   }
@@ -410,7 +408,6 @@ class GraphRuntimeCodegen : public backend::MemoizedExprTranslator<std::vector<G
     //   LOG(FATAL) << "TVM only support calls to primitive functions "
     //              << "(i.e functions composed of fusable operator invocations)";
     // }
-    std::cout << PrettyPrint(call) << std::endl;
 
     if (auto op_node = call->op.as<OpNode>()) {
       if (op_node->name != "prim_fn_call") {

diff --git a/src/relay/backend/tir_compiler.cc b/src/relay/backend/tir_compiler.cc
@@ -887,17 +887,9 @@ class LowerTensorExpr : public ExprMutator {
     }
 
     // Process inputs.
-    bool skip_first;
     Array<Expr> args;
-    for (auto arg : expr->args) {
-      // The first input is a function, not a tensor.
-      if (skip_first) {
-        skip_first = false;
-        args.push_back(arg);
-        continue;
-      }
-
-      args.push_back(VisitExpr(arg));
+    for (size_t i = 0; i < expr->args.size(); i++) {
+      args.push_back(VisitExpr(expr->args[i]));
     }
 
     Target target;