tkonolige commented on code in PR #11358:
URL: https://github.com/apache/tvm/pull/11358#discussion_r967278430
##########
src/runtime/vm/vm.cc:
##########
@@ -518,13 +584,54 @@ int64_t VirtualMachine::LoadScalarInt(Index r) const {
return result;
}
-void VirtualMachine::RunLoop() {
+Index VirtualMachine::GetResultRegisterIndex() const {
+ Index op_index = 0;
+ while (code_[op_index].op != Opcode::Ret) {
+ ++op_index;
+ }
+
+ return code_[op_index].result;
+}
+
+void VirtualMachine::CalculatePreResultOpIndex(Index res_index) {
+ if (preresult_op_index_ == -1) {
+ preresult_op_index_ = 0;
+ while (code_[preresult_op_index_].dst != res_index) {
+ ++preresult_op_index_;
+ }
+ }
+}
+
+void VirtualMachine::CollectOutputTensorRegIndices(const std::string&
func_name) {
+ if (!output_tensor_reg_indices_[func_name].empty()) {
+ return;
+ }
+
+ auto& reg_indices = output_tensor_reg_indices_[func_name];
+ Index res_index = GetResultRegisterIndex();
+ CalculatePreResultOpIndex(res_index);
+ auto& preres_instr = code_[preresult_op_index_];
+ auto op_code = preres_instr.op;
+ if (op_code == Opcode::AllocTensor) {
+ reg_indices.emplace_back(res_index);
+ } else if (op_code == Opcode::AllocADT) {
+ for (Index i = 0; i < preres_instr.num_fields; ++i) {
+ reg_indices.push_back(preres_instr.datatype_fields[i]);
+ }
+ } else if (op_code == Opcode::ReshapeTensor) {
+ reg_indices.push_back(preres_instr.reshape_tensor.tensor);
+ } else {
+ LOG(WARNING) << "Operation " << size_t(op_code) << " is not supported for
set_outputs method";
+ }
+}
+
+void VirtualMachine::RunLoop(const std::vector<Index>&
output_tensor_reg_indices) {
ICHECK(this->exec_);
ICHECK(this->code_);
pc_ = 0;
Index frame_start = frames_.size();
- while (true) {
- main_loop:
+ bool iterate = true;
+ while (iterate) {
Review Comment:
If you want to change it you can submit a separate PR to do so. I don't have
enough knowlege to say if there was a good reason behind 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]
