masahi commented on code in PR #15419:
URL: https://github.com/apache/tvm/pull/15419#discussion_r1282880516
##########
src/runtime/vm/bytecode.cc:
##########
@@ -338,14 +340,22 @@ Instruction Instruction::AllocTensorReg(RegName storage,
RegName offset, RegName
}
Instruction Instruction::AllocStorage(RegName size, Index alignment,
DLDataType dtype_hint,
- Index device_index, RegName dst) {
+ Index device_index, const
std::vector<int64_t>& shape,
+ RegName dst) {
Instruction instr;
instr.op = Opcode::AllocStorage;
instr.dst = dst;
instr.alloc_storage.allocation_size = size;
instr.alloc_storage.alignment = alignment;
instr.alloc_storage.dtype_hint = dtype_hint;
instr.alloc_storage.device_index = device_index;
+ instr.alloc_storage.ndim = static_cast<uint32_t>(shape.size());
+ if (instr.alloc_storage.ndim > 0) {
+ instr.alloc_storage.shape = new int64_t[shape.size()];
Review Comment:
This alloc is leaking?
##########
src/runtime/vm/vm.cc:
##########
@@ -819,17 +820,37 @@ void VirtualMachine::RunLoop(const std::vector<Index>&
output_tensor_reg_indices
}
case Opcode::AllocStorage: {
OpStartHook(instr);
- auto size = LoadScalarInt(instr.alloc_storage.allocation_size);
- auto alignment = instr.alloc_storage.alignment;
auto storage_obj = SimpleObjAllocator().make_object<StorageObj>();
Allocator* allocator = GetAllocator(instr.alloc_storage.device_index);
ICHECK(allocator) << "Did you forget to init the VirtualMachine with
devices?";
- VLOG(2) << "allocating with allocation_size=" << size << ",
alignment=" << alignment
- << ", dtype_hint=" <<
DLDataType2String(instr.alloc_storage.dtype_hint)
- << ", device_index=" << instr.alloc_storage.device_index;
-
- storage_obj->buffer = allocator->Alloc(size, alignment,
instr.alloc_storage.dtype_hint);
+ std::string mem_scope =
exec_->virtual_devices[instr.alloc_storage.device_index].second;
Review Comment:
Put this decl under `if (instr.alloc_storage.ndim > 0) {`
`const std::string& mem_scope = ...`
##########
src/runtime/vm/vm.cc:
##########
@@ -899,8 +920,9 @@ void VirtualMachine::RunLoop(const std::vector<Index>&
output_tensor_reg_indices
ICHECK_EQ(actual_src_dev.device_type, inst_src_dev.device_type);
ICHECK_EQ(actual_src_dev.device_id, inst_src_dev.device_id);
Device dst_dev = GetDevice(instr.device_copy.dst_device_index);
+ auto mem_scope =
exec_->virtual_devices[instr.device_copy.dst_device_index].second;
- NDArray dst_data = src_data.CopyTo(dst_dev);
+ NDArray dst_data = src_data.CopyTo(dst_dev, String(mem_scope));
Review Comment:
This is creating a string copy. Possible to avoid it?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
