From 2e0c5510393d526b28c7e63a85b5f1a906820dbb Mon Sep 17 00:00:00 2001
From: Chris Sullivan <csullivan@octoml.ai>
Date: Fri, 29 Jan 2021 15:25:34 -0800
Subject: [PATCH] Fix a few bugs: 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

---
 python/tvm/relay/backend/_backend.py       |  1 -
 src/relay/backend/graph_runtime_codegen.cc |  3 ---
 src/relay/backend/tir_compiler.cc          | 12 ++----------
 3 files changed, 2 insertions(+), 14 deletions(-)

diff --git a/python/tvm/relay/backend/_backend.py b/python/tvm/relay/backend/_backend.py
index 0a6232bf7f58..65b0c0ba87c7 100644
--- 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
index f47603d7352f..416540db2286 100644
--- 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
index 1de9cb7ca64e..5df64c024a66 100644
--- 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;