diff --git a/include/tvm/runtime/vm.h b/include/tvm/runtime/vm.h index 6cb568822e6c..12ba7eda3b30 100644 --- 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 index 5fb4909b57c7..e05befcc70dd 100644 --- 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 index 7f45ec92a630..a23a7bce4ef5 100644 --- 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 index 9cd27050e54c..30c03391bc3c 100644 --- 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(instr.fields[4])); @@ -580,7 +580,7 @@ Instruction DeserializeInstruction(const VMInstructionSerializer& instr) { Index ndim = instr.fields[5]; RegName dst = instr.fields[6]; - std::vector shape = ExtractFields(instr.fields, 6, ndim); + std::vector shape = ExtractFields(instr.fields, 7, ndim); return Instruction::AllocTensor(storage_reg, offset, shape, dtype, dst); }