alloc_storage and rebase

python/tvm/relay/transform/memory_alloc.py

@@ -174,7 +174,7 @@ def dynamic_invoke(self, scope, func, ins, new_args, out_types, ret_type):
             size = self.compute_storage_in_relay(
                 out_shape, out_type.dtype)
             alignment = self.compute_alignment(out_type.dtype)
-            sto = scope.let("storage_{i}".format(i=i), self.alloc_storage(
+            sto = scope.let("storage_{i}".format(i=i), alloc_storage(
                 size, alignment, self.default_context, out_type.dtype))
             storages.append(sto)
 

src/relay/backend/vm/compiler.cc

@@ -894,13 +894,6 @@ transform::Sequential MemoryOpt(tvm::Target host_target) {
   // Fuse the shape functions.
   pass_seqs.push_back(transform::FuseOps());
 
-  // Inline the functions that are lifted to the module scope. We perform this
-  // pass after all other optimization passes but before the memory allocation
-  // pass. This is because memory allocation pass will insert `invoke_tvm_op`
-  // and we use these ops to invoke the symbols in the module generated by
-  // external codegen.
-  pass_seqs.push_back(transform::Inline());
-
   // Manifest the allocations needed for the shape functions.
   pass_seqs.push_back(transform::ManifestAlloc(host_target));
 
@@ -963,6 +956,13 @@ IRModule VMCompiler::OptimizeModule(const IRModule& mod, const TargetsMap& targe
   pass_seqs.push_back(transform::LambdaLift());
   pass_seqs.push_back(transform::InlinePrimitives());
 
+  // Inline the functions that are lifted to the module scope. We perform this
+  // pass after all other optimization passes but before the memory allocation
+  // pass. This is because memory allocation pass will insert `invoke_tvm_op`
+  // and we use these ops to invoke the symbols in the module generated by
+  // external codegen.
+  pass_seqs.push_back(transform::Inline());
+
   // Memory optimization
   pass_seqs.push_back(MemoryOpt(this->target_host_));