move TECompiler into VMCompilerContext

apache · Jul 20, 2021 · 8e3d7ff · 8e3d7ff
1 parent c4d08d2
commit 8e3d7ff
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc
@@ -466,7 +466,7 @@ class VMFunctionCompiler : ExprFunctor<void(const Expr& expr)> {
   void EmitShapeFunc(Function func, Array<Expr> inputs, Array<Expr> outputs) {
     // Lower shape function
     CCacheKey key(func, target_host_);
-    auto cfunc = compiler_->LowerShapeFunc(key);
+    auto cfunc = context_->compiler->LowerShapeFunc(key);
     int op_index = -1;
     // pick the only function inside the context
     ICHECK_EQ(cfunc->funcs->functions.size(), 1);
@@ -552,7 +552,7 @@ class VMFunctionCompiler : ExprFunctor<void(const Expr& expr)> {
 
     CCacheKey key(func, target);
     auto mangle_fn = [](String name) { return name; };
-    auto cfunc = compiler_->Lower(key, mangle_fn);
+    auto cfunc = context_->compiler->Lower(key, mangle_fn);
 
     auto op_index = -1;
     if (func->GetAttr<String>(attr::kCompiler).defined()) {
@@ -858,8 +858,6 @@ class VMFunctionCompiler : ExprFunctor<void(const Expr& expr)> {
   size_t last_register_;
   /*! \brief Total number of virtual registers allocated. */
   size_t registers_num_;
-  /*! \brief Compiler engine to lower primitive functions. */
-  TECompiler compiler_;
   /*! \brief Global shared meta data */
   VMCompilerContext* context_;
   /*! \brief Target devices. */
@@ -1185,8 +1183,7 @@ void VMCompiler::Codegen() {
     }
   }
 
-  TECompiler compiler;
-  auto ext_mods = compiler->LowerExternalFunctions();
+  auto ext_mods = context_.compiler->LowerExternalFunctions();
 
   runtime::Module lib;
   if (funcs.size() > 0) {

diff --git a/src/relay/backend/vm/compiler.h b/src/relay/backend/vm/compiler.h
@@ -76,6 +76,8 @@ struct VMCompilerContext {
   TagMap tag_map;
   // Map from global var to a unique integer
   GlobalMap global_map;
+  // TECompiler to do lowering
+  tec::TECompiler compiler;
   // List of constants
   std::vector<NDArray> constants;
   // Device type for constants