Merge pull request #9 from zhiics/coalesce-storage

some fixes
jroesch · Apr 22, 2020 · 36da8f2 · 36da8f2
2 parents e4b1391 + 439d59d
commit 36da8f2
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 8 deletions.
diff --git a/include/tvm/runtime/vm.h b/include/tvm/runtime/vm.h
@@ -270,6 +270,7 @@ struct Instruction {
   /*!
    * \brief Construct an allocate tensor instruction with constant shape.
    * \param storage The storage to allocate out of.
+   * \param offset The offset into the storage to allocate from.
    * \param shape The shape of the tensor.
    * \param dtype The dtype of the tensor.
    * \param dst The destination register.

diff --git a/python/tvm/relay/transform/memory_alloc.py b/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)
 

diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc
@@ -906,8 +906,6 @@ transform::Sequential MemoryOpt(tvm::Target host_target) {
 
   // Perform memory planning in order to coalesce/reduce allocations.
   pass_seqs.push_back(transform::MemoryPlan());
-  // Compute away possibly introduced constant computation.
-  pass_seqs.push_back(transform::FoldConstant());
 
   return transform::Sequential(pass_seqs);
 }
@@ -965,8 +963,7 @@ IRModule VMCompiler::OptimizeModule(const IRModule& mod, const TargetsMap& targe
   pass_seqs.push_back(transform::LambdaLift());
   pass_seqs.push_back(transform::InlinePrimitives());
 
-
-
+  // Memory optimization
   pass_seqs.push_back(MemoryOpt(this->target_host_));
 
   transform::Sequential seq(pass_seqs);

diff --git a/src/runtime/vm/executable.cc b/src/runtime/vm/executable.cc
@@ -314,7 +314,7 @@ VMInstructionSerializer SerializeInstruction(const Instruction& instr) {
       break;
     }
     case Opcode::AllocTensor: {
-      // Number of fields = 6 + instr.alloc_tensor.ndim
+      // Number of fields = 7 + instr.alloc_tensor.ndim
       fields.push_back(instr.alloc_tensor.storage);
       fields.push_back(instr.alloc_tensor.offset);
       // Save `DLDataType` and the dst register.
@@ -565,7 +565,7 @@ Instruction DeserializeInstruction(const VMInstructionSerializer& instr) {
       return Instruction::InvokePacked(packed_index, arity, output_size, args);
     }
     case Opcode::AllocTensor: {
-      // Number of fields = 6 + instr.alloc_tensor.ndim
+      // Number of fields = 7 + instr.alloc_tensor.ndim
       DCHECK_GE(instr.fields.size(), 7U);
       DCHECK_EQ(instr.fields.size(), 7U + static_cast<size_t>(instr.fields[4]));
 
@@ -580,7 +580,7 @@ Instruction DeserializeInstruction(const VMInstructionSerializer& instr) {
       Index ndim = instr.fields[5];
       RegName dst = instr.fields[6];
 
-      std::vector<Index> shape = ExtractFields(instr.fields, 6, ndim);
+      std::vector<Index> shape = ExtractFields(instr.fields, 7, ndim);
 
       return Instruction::AllocTensor(storage_reg, offset, shape, dtype, dst);
     }