vvchernov commented on code in PR #11358:
URL: https://github.com/apache/tvm/pull/11358#discussion_r966808686
##########
src/runtime/vm/vm.cc:
##########
@@ -810,21 +906,90 @@ void VirtualMachine::RunLoop() {
WriteRegister(instr.dst, dst_data);
OpStopHook();
pc_++;
- goto main_loop;
+ break;
}
case Opcode::KillRegister: {
OpStartHook(instr);
WriteRegister(instr.dst, ObjectRef());
OpStopHook();
pc_++;
- goto main_loop;
+ break;
}
default:
LOG(FATAL) << "Unknown instruction opcode: " << int(instr.op);
}
}
}
+void VirtualMachine::WriteAllocatedTensor(const Instruction& instr) {
+ auto shape = std::vector<int64_t>(instr.alloc_tensor.ndim);
+
+ for (uint32_t i = 0; i < instr.alloc_tensor.ndim; ++i) {
+ shape[i] = instr.alloc_tensor.shape[i];
+ }
+
+ auto storage_obj = ReadRegister(instr.alloc_tensor.storage);
+ auto offset = LoadScalarInt(instr.alloc_tensor.offset);
+ auto storage = Downcast<Storage>(storage_obj);
+ auto obj = storage->AllocNDArray(offset, shape, instr.alloc_tensor.dtype);
+ VLOG(2) << "allocated "
+ << RuntimeObject2String(obj, GetDevice(exec_->host_device_index),
+ /*show_contents=*/false);
+
+ WriteRegister(instr.dst, obj);
+}
+
+void VirtualMachine::WriteAllocatedTensorFromOutside(const Instruction& instr)
{
+ // External tensor(s) has been already written to the register (instr.dst)
+ auto ex_arr = Downcast<NDArray>(ReadRegister(instr.dst));
+ auto ex_shape = ex_arr.Shape();
+ auto ex_size = ex_shape.size();
+ auto ex_dtype = ex_arr->dtype;
+
+ auto in_size = instr.alloc_tensor.ndim;
+ auto in_dtype = instr.alloc_tensor.dtype;
+ ICHECK_EQ(TypeEqual(in_dtype, ex_dtype), true)
+ << "Data types mismatching for internal and external output tensors";
+
+ bool size_check = false;
+ if (ex_size != in_size) {
+ size_check = true;
+ } else {
+ for (size_t i = 0; i < in_size; ++i) {
+ if (ex_shape[i] != instr.alloc_tensor.shape[i]) {
+ size_check = true;
+ break;
+ }
+ }
+ }
+
+ if (size_check) {
+ // Match element number
+ size_t in_el_num = 1, ex_el_num = 1;
+ for (size_t i = 0; i < ex_size; ++i) {
+ ex_el_num *= ex_shape[i];
+ }
+ for (size_t i = 0; i < in_size; ++i) {
+ in_el_num *= instr.alloc_tensor.shape[i];
+ }
+ ICHECK_EQ(in_el_num, ex_el_num)
+ << "Element number mismatching of internal and external output
tensors";
+ if (code_[preresult_op_index_].op == Opcode::ReshapeTensor) {
+ int64_t* dims = instr.alloc_tensor.shape;
+ std::vector<int64_t> ref_shape(dims, dims + int64_t(in_size));
+ auto reshaped_tensor = ex_arr.CreateView(ref_shape, ex_dtype);
+ WriteRegister(instr.dst, reshaped_tensor);
+ } else {
+ LOG_ERROR << "Internal and external output tensor shapes are mismatched";
Review Comment:
Thanks! fixed
--
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]
