You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2022/05/12 11:30:20 UTC

[GitHub] [tvm] ruxiliang opened a new pull request, #11286: Dynamic management

ruxiliang opened a new pull request, #11286:
URL: https://github.com/apache/tvm/pull/11286

   Thanks for contributing to TVM!   Please refer to guideline https://tvm.apache.org/docs/contribute/ for useful information and tips. After the pull request is submitted, please request code reviews from [Reviewers](https://github.com/apache/incubator-tvm/blob/master/CONTRIBUTORS.md#reviewers) by @ them in the pull request thread.
   
   This pull request adds the [checkpointing](https://github.com/cybertronai/gradient-checkpointing) feature for the GraphExecutor for TVM. Adding this will help implement heuristic strategies like @MarisaKirisame 's  [dynamic tensor materialization](https://arxiv.org/abs/2006.09616) which could help model reducing model memory usage.


-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] altanh commented on a diff in pull request #11286: Dynamic management

Posted by GitBox <gi...@apache.org>.
altanh commented on code in PR #11286:
URL: https://github.com/apache/tvm/pull/11286#discussion_r877531120


##########
include/tvm/runtime/ndarray.h:
##########
@@ -73,7 +73,7 @@ class NDArray : public ObjectRef {
   explicit NDArray(ObjectPtr<Object> data) : ObjectRef(data) {}
 
   /*! \brief reset the content of NDArray to be nullptr */
-  inline void reset();
+  inline void reset(){ data_.reset(); }

Review Comment:
   is this change necessary?



##########
src/runtime/graph_executor/graph_executor.cc:
##########
@@ -59,10 +61,98 @@ inline size_t GetDataAlignment(const DLTensor& arr) {
 void GraphExecutor::Run() {
   // setup the array and requirements.
   for (size_t i = 0; i < op_execs_.size(); ++i) {
-    if (op_execs_[i]) op_execs_[i]();
+    if (op_execs_[i]) {
+      PerformOp(i);
+    }
+  }
+}
+
+void GraphExecutor::RematerializeTensor(size_t nid) { PerformOp(nid); }
+
+void GraphExecutor::PerformOp(size_t nid) {
+  auto node = nodes_[nid];
+  // first we need to get all the input tensors and add the reference count
+  // if they are not in memory, rematerialize them
+  for (const auto& input : node.inputs) {
+    // also we do not care about cpu
+    if (pool_entry_[get_sid(entry_id(input))].device_type == (int)(kDLCPU)) continue;
+    // we do not need to care about input and parameter
+    if (std::find(input_nodes_.begin(), input_nodes_.end(), entry_id(input)) != input_nodes_.end())
+      continue;
+    if (std::find(nodes_not_evicted_.begin(), nodes_not_evicted_.end(), entry_id(input)) != nodes_not_evicted_.end())
+      continue;
+    // otherwise we should rematerialize the input tensors that are not in the memory
+    // and update the reference count at the same time
+
+    if (is_evicted(entry_id(input))) {
+        RematerializeTensor(entry_id(input));
+    }
+    ref_cnt[entry_id(input)] += 1;
+  }
+  // now all of the inputs are recovered, we can set up to compute the location of input/output
+  // compute current input memory location of the tensor
+  std::unordered_set<uint32_t> input_node_eids;
+  for (unsigned int nid_ : input_nodes_) {
+    input_node_eids.insert(entry_id(nid_, 0));
+  }
+  std::unordered_set<uint32_t> output_node_eids;
+  for (auto& output : outputs_) {
+    output_node_eids.insert(entry_id(output));
+  }
+
+  const auto& inode = nodes_[nid];
+  std::vector<DLTensor> args;
+
+  for (const auto& e : inode.inputs) {
+    uint32_t eid = this->entry_id(e);
+    args.push_back(*(data_entry_[eid].operator->()));
+  }
+
+  /*
+   * Note that some output tensors might not be allocated because some op might produce plural
+   * values Thus we should check if the outputs were allocated and set up the data entry as well
+   */
+  for (uint32_t index = 0; index < inode.param.num_outputs; ++index) {
+    uint32_t eid = this->entry_id(nid, index);
+    set_tensor(eid);
+    args.push_back(*(data_entry_[eid].operator->()));
+  }
+
+  // Now we should update the memory location in TVM Operator
+  std::shared_ptr<OpArgs> updated_args = UpdateTVMOp(nid, nodes_[nid].param, args);
+
+  UpdateInputOutputTensors(input_node_eids, output_node_eids, nid, updated_args);
+
+  // now we can execute the operator
+  op_execs_[nid]();
+  // at last, we should perform eviction
+  PerformEviction(nid);
+}
+
+void GraphExecutor::PerformEviction(size_t nid) {
+  auto node = nodes_[nid];
+  for (const auto& input : node.inputs) {
+    // we do not care about cpu

Review Comment:
   I understand the reasoning, but should this be hardcoded?



##########
src/runtime/graph_executor/graph_executor.cc:
##########
@@ -422,14 +527,31 @@ void GraphExecutor::SetupStorage() {
   for (size_t i = 0; i < data_entry_.size(); ++i) {
     int storage_id = attrs_.storage_id[i];
     ICHECK_LT(static_cast<size_t>(storage_id), storage_pool_.size());
-    data_entry_[i] = storage_pool_[storage_id].CreateView(attrs_.shape[i], vtype[i]);
+    // just make view for allocated parameters and tensors in cpu
+    if (pool_entry_[storage_id].device_type == static_cast<int>(kDLCPU)){
+      data_entry_[i] = storage_pool_[storage_id].CreateView(attrs_.shape[i], vtype[i]);
+    }
 
-    const DLTensor* tmp = data_entry_[i].operator->();
-    data_alignment_[i] = details::GetDataAlignment(*tmp);
+    // determine memory usage using grabbed vtype
+    size_t align = (vtype[i].bits / 8) * vtype[i].lanes;
+    data_alignment_[i] = align < kAllocAlignment ? kAllocAlignment : align;
   }
 }
 
 void GraphExecutor::SetupOpExecs() {
+  ref_cnt.resize(num_node_entries(),0);
+  nodes_not_evicted_ = {};
+  for(size_t nid = 0;nid < num_node_entries();nid++){
+    auto slice_node = nodes_[nid];
+    if (slice_node.name.find("slice") != std::string::npos){

Review Comment:
   can you explain this?



##########
src/runtime/graph_executor/graph_executor.cc:
##########
@@ -59,10 +61,98 @@ inline size_t GetDataAlignment(const DLTensor& arr) {
 void GraphExecutor::Run() {
   // setup the array and requirements.
   for (size_t i = 0; i < op_execs_.size(); ++i) {
-    if (op_execs_[i]) op_execs_[i]();
+    if (op_execs_[i]) {
+      PerformOp(i);
+    }
+  }
+}
+
+void GraphExecutor::RematerializeTensor(size_t nid) { PerformOp(nid); }
+
+void GraphExecutor::PerformOp(size_t nid) {
+  auto node = nodes_[nid];
+  // first we need to get all the input tensors and add the reference count
+  // if they are not in memory, rematerialize them
+  for (const auto& input : node.inputs) {
+    // also we do not care about cpu
+    if (pool_entry_[get_sid(entry_id(input))].device_type == (int)(kDLCPU)) continue;
+    // we do not need to care about input and parameter
+    if (std::find(input_nodes_.begin(), input_nodes_.end(), entry_id(input)) != input_nodes_.end())
+      continue;
+    if (std::find(nodes_not_evicted_.begin(), nodes_not_evicted_.end(), entry_id(input)) != nodes_not_evicted_.end())
+      continue;
+    // otherwise we should rematerialize the input tensors that are not in the memory
+    // and update the reference count at the same time
+
+    if (is_evicted(entry_id(input))) {
+        RematerializeTensor(entry_id(input));
+    }
+    ref_cnt[entry_id(input)] += 1;
+  }
+  // now all of the inputs are recovered, we can set up to compute the location of input/output
+  // compute current input memory location of the tensor
+  std::unordered_set<uint32_t> input_node_eids;
+  for (unsigned int nid_ : input_nodes_) {
+    input_node_eids.insert(entry_id(nid_, 0));
+  }
+  std::unordered_set<uint32_t> output_node_eids;
+  for (auto& output : outputs_) {
+    output_node_eids.insert(entry_id(output));
+  }
+
+  const auto& inode = nodes_[nid];
+  std::vector<DLTensor> args;
+
+  for (const auto& e : inode.inputs) {
+    uint32_t eid = this->entry_id(e);
+    args.push_back(*(data_entry_[eid].operator->()));
+  }
+
+  /*
+   * Note that some output tensors might not be allocated because some op might produce plural
+   * values Thus we should check if the outputs were allocated and set up the data entry as well
+   */
+  for (uint32_t index = 0; index < inode.param.num_outputs; ++index) {
+    uint32_t eid = this->entry_id(nid, index);
+    set_tensor(eid);
+    args.push_back(*(data_entry_[eid].operator->()));
+  }
+
+  // Now we should update the memory location in TVM Operator
+  std::shared_ptr<OpArgs> updated_args = UpdateTVMOp(nid, nodes_[nid].param, args);
+
+  UpdateInputOutputTensors(input_node_eids, output_node_eids, nid, updated_args);
+
+  // now we can execute the operator
+  op_execs_[nid]();
+  // at last, we should perform eviction
+  PerformEviction(nid);
+}
+
+void GraphExecutor::PerformEviction(size_t nid) {

Review Comment:
   if I understand correctly, this function decreases the ref count and tries to eagerly evict the inputs of `nid`. I would prefer if the logic was split out so that the reference decrement + eager eviction applied directly to `nid` (and then call this function in a loop over the inputs once an op has been performed).
   
   would also suggest a different name, maybe call it `ReleaseTensor` or something similar, since this is different from the core eviction loop in DTR



##########
src/runtime/graph_executor/graph_executor.cc:
##########
@@ -59,10 +61,98 @@ inline size_t GetDataAlignment(const DLTensor& arr) {
 void GraphExecutor::Run() {
   // setup the array and requirements.
   for (size_t i = 0; i < op_execs_.size(); ++i) {
-    if (op_execs_[i]) op_execs_[i]();
+    if (op_execs_[i]) {
+      PerformOp(i);
+    }
+  }
+}
+
+void GraphExecutor::RematerializeTensor(size_t nid) { PerformOp(nid); }

Review Comment:
   could you subclass the `GraphExecutor` so that all these changes don't become default behavior? hopefully it's not too tricky to just override the relevant methods



-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] slyubomirsky commented on a diff in pull request #11286: Dynamic management

Posted by GitBox <gi...@apache.org>.
slyubomirsky commented on code in PR #11286:
URL: https://github.com/apache/tvm/pull/11286#discussion_r881196061


##########
src/runtime/graph_executor/graph_executor.cc:
##########
@@ -448,74 +570,65 @@ void GraphExecutor::SetupOpExecs() {
   for (uint32_t nid = 0; nid < this->GetNumOfNodes(); ++nid) {
     const auto& inode = nodes_[nid];
     if (inode.op_type == "null") continue;
-    std::vector<DLTensor> args;
-    for (const auto& e : inode.inputs) {
-      uint32_t eid = this->entry_id(e);
-      args.push_back(*(data_entry_[eid].operator->()));
-    }
-    for (uint32_t index = 0; index < inode.param.num_outputs; ++index) {
-      uint32_t eid = this->entry_id(nid, index);
-      args.push_back(*(data_entry_[eid].operator->()));
-    }
-    ICHECK(inode.op_type == "tvm_op") << "Can only take tvm_op as op";
+    std::vector<DLTensor> args = {};
 
     std::shared_ptr<OpArgs> op_args = nullptr;
-    std::tie(op_execs_[nid], op_args) = CreateTVMOp(inode.param, args);
-
-    for (size_t i = 0; i < inode.inputs.size(); i++) {
-      uint32_t input_eid = this->entry_id(inode.inputs[i]);
-      // check if op input is model input
-      if (input_node_eids.count(input_eid) > 0) {
-        input_dltensors_[input_eid].push_back(
-            static_cast<DLTensor*>(op_args->arg_values[i].v_handle));
+    std::tie(op_execs_[nid], op_args) = CreateTVMOp(nid, inode.param, args);
+  }
+}
+
+void GraphExecutor::UpdateInputOutputTensors(const std::unordered_set<uint32_t>& input_node_eids,
+                                             const std::unordered_set<uint32_t>& output_node_eids,
+                                             uint32_t nid,
+                                             std::shared_ptr<GraphExecutor::OpArgs> op_args) {
+  const auto& inode = nodes_[nid];
+  for (size_t i = 0; i < inode.inputs.size(); i++) {
+    uint32_t input_eid = this->entry_id(inode.inputs[i]);
+    // check if op input is model input
+    if (input_node_eids.count(input_eid) > 0) {
+      auto& eid_input = input_dltensors_[input_eid];
+      auto input_tensor = static_cast<DLTensor*>(op_args->arg_values[i].v_handle);
+      if (std::find(eid_input.begin(), eid_input.end(), input_tensor) == eid_input.end()) {
+        eid_input.push_back(input_tensor);
       }
-      // check if any model output is the input of the op
-      if (output_node_eids.count(input_eid) > 0) {
-        both_output_opinput_dltensors_[input_eid].push_back(
-            static_cast<DLTensor*>(op_args->arg_values[i].v_handle));
+    }
+    // check if any model output is the input of the op
+    if (output_node_eids.count(input_eid) > 0) {
+      auto& eid_both = both_output_opinput_dltensors_[input_eid];
+      auto both_tensor = static_cast<DLTensor*>(op_args->arg_values[i].v_handle);
+      if(std::find(eid_both.begin(), eid_both.end(), both_tensor) == eid_both.end()){
+        eid_both.push_back(both_tensor);
       }
     }
+  }
 
-    for (uint32_t i = inode.inputs.size(); i < inode.inputs.size() + inode.param.num_outputs; ++i) {
-      uint32_t output_eid = this->entry_id(nid, i - inode.inputs.size());
-      // check if op output is model output
-      if (output_node_eids.count(output_eid) > 0) {
-        output_dltensors_[output_eid].push_back(
-            static_cast<DLTensor*>(op_args->arg_values[i].v_handle));
+  for (uint32_t i = inode.inputs.size(); i < inode.inputs.size() + inode.param.num_outputs; ++i) {
+    uint32_t output_eid = this->entry_id(nid, i - inode.inputs.size());
+    // check if op output is model output
+    if (output_node_eids.count(output_eid) > 0) {
+      auto& eid_out = output_dltensors_[output_eid];
+      auto out_tensor = static_cast<DLTensor*>(op_args->arg_values[i].v_handle);
+      if(std::find(eid_out.begin(), eid_out.end(), out_tensor) == eid_out.end()){
+        eid_out.push_back(out_tensor);
       }
     }
   }
 }
-
 std::pair<std::function<void()>, std::shared_ptr<GraphExecutor::OpArgs> >
-GraphExecutor::CreateTVMOp(const TVMOpParam& param, const std::vector<DLTensor>& args) {
+GraphExecutor::CreateTVMOp(uint32_t nid, const TVMOpParam& param,
+                           const std::vector<DLTensor>& args) {
   std::shared_ptr<GraphExecutor::OpArgs> arg_ptr = std::make_shared<GraphExecutor::OpArgs>();
-  // setup address.
-  arg_ptr->args = args;
-  if (param.flatten_data) {
-    arg_ptr->shape_data.resize(arg_ptr->args.size());
-  }
-  for (size_t i = 0; i < arg_ptr->args.size(); ++i) {
-    TVMValue v;
-    DLTensor* t = &arg_ptr->args[i];
-    v.v_handle = t;
-    arg_ptr->arg_values.push_back(v);
-    arg_ptr->arg_tcodes.push_back(kTVMDLTensorHandle);
-    if (param.flatten_data) {
-      arg_ptr->shape_data[i] =
-          std::accumulate(t->shape, t->shape + t->ndim, 1, std::multiplies<int64_t>());
-      t->ndim = 1;
-      t->shape = &(arg_ptr->shape_data[i]);
-    }
-  }
+  // set up the memory location map
+  op_mem_location_map_[nid] = arg_ptr;
 
   if (param.func_name == "__nop") {
     return {[]() {}, arg_ptr};
   } else if (param.func_name == "__copy") {
     // Perform cross device data copy.
     // Directly copy data from the input to the output.
     // TODO(mbs): device_copy cleanup.
-    auto fexec = [arg_ptr]() {
+    auto fexec = [nid, this]() {
+      auto arg_ptr = this->op_mem_location_map_[nid];

Review Comment:
   Perhaps it might be more economical to define `arg_ptr` outside the lambda and capture only that? (Same with the other lambdas that capture `this`.) I doubt it's a serious issue but it may be worth checking.



-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] ruxiliang commented on a diff in pull request #11286: Dynamic management

Posted by GitBox <gi...@apache.org>.
ruxiliang commented on code in PR #11286:
URL: https://github.com/apache/tvm/pull/11286#discussion_r877682467


##########
include/tvm/runtime/ndarray.h:
##########
@@ -73,7 +73,7 @@ class NDArray : public ObjectRef {
   explicit NDArray(ObjectPtr<Object> data) : ObjectRef(data) {}
 
   /*! \brief reset the content of NDArray to be nullptr */
-  inline void reset();
+  inline void reset(){ data_.reset(); }

Review Comment:
   I remembered that if I do not add this the ld link cannot find for the reset method. I will check this again.



##########
src/runtime/graph_executor/graph_executor.cc:
##########
@@ -59,10 +61,98 @@ inline size_t GetDataAlignment(const DLTensor& arr) {
 void GraphExecutor::Run() {
   // setup the array and requirements.
   for (size_t i = 0; i < op_execs_.size(); ++i) {
-    if (op_execs_[i]) op_execs_[i]();
+    if (op_execs_[i]) {
+      PerformOp(i);
+    }
+  }
+}
+
+void GraphExecutor::RematerializeTensor(size_t nid) { PerformOp(nid); }
+
+void GraphExecutor::PerformOp(size_t nid) {
+  auto node = nodes_[nid];
+  // first we need to get all the input tensors and add the reference count
+  // if they are not in memory, rematerialize them
+  for (const auto& input : node.inputs) {
+    // also we do not care about cpu
+    if (pool_entry_[get_sid(entry_id(input))].device_type == (int)(kDLCPU)) continue;
+    // we do not need to care about input and parameter
+    if (std::find(input_nodes_.begin(), input_nodes_.end(), entry_id(input)) != input_nodes_.end())
+      continue;
+    if (std::find(nodes_not_evicted_.begin(), nodes_not_evicted_.end(), entry_id(input)) != nodes_not_evicted_.end())
+      continue;
+    // otherwise we should rematerialize the input tensors that are not in the memory
+    // and update the reference count at the same time
+
+    if (is_evicted(entry_id(input))) {
+        RematerializeTensor(entry_id(input));
+    }
+    ref_cnt[entry_id(input)] += 1;
+  }
+  // now all of the inputs are recovered, we can set up to compute the location of input/output
+  // compute current input memory location of the tensor
+  std::unordered_set<uint32_t> input_node_eids;
+  for (unsigned int nid_ : input_nodes_) {
+    input_node_eids.insert(entry_id(nid_, 0));
+  }
+  std::unordered_set<uint32_t> output_node_eids;
+  for (auto& output : outputs_) {
+    output_node_eids.insert(entry_id(output));
+  }
+
+  const auto& inode = nodes_[nid];
+  std::vector<DLTensor> args;
+
+  for (const auto& e : inode.inputs) {
+    uint32_t eid = this->entry_id(e);
+    args.push_back(*(data_entry_[eid].operator->()));
+  }
+
+  /*
+   * Note that some output tensors might not be allocated because some op might produce plural
+   * values Thus we should check if the outputs were allocated and set up the data entry as well
+   */
+  for (uint32_t index = 0; index < inode.param.num_outputs; ++index) {
+    uint32_t eid = this->entry_id(nid, index);
+    set_tensor(eid);
+    args.push_back(*(data_entry_[eid].operator->()));
+  }
+
+  // Now we should update the memory location in TVM Operator
+  std::shared_ptr<OpArgs> updated_args = UpdateTVMOp(nid, nodes_[nid].param, args);
+
+  UpdateInputOutputTensors(input_node_eids, output_node_eids, nid, updated_args);
+
+  // now we can execute the operator
+  op_execs_[nid]();
+  // at last, we should perform eviction
+  PerformEviction(nid);
+}
+
+void GraphExecutor::PerformEviction(size_t nid) {
+  auto node = nodes_[nid];
+  for (const auto& input : node.inputs) {
+    // we do not care about cpu

Review Comment:
   I think leaving host tensors alone is not necessary, but to get the result of the experiments quickly I ignored those tensors so that I can get the results quickly. 



##########
src/runtime/graph_executor/graph_executor.cc:
##########
@@ -59,10 +61,98 @@ inline size_t GetDataAlignment(const DLTensor& arr) {
 void GraphExecutor::Run() {
   // setup the array and requirements.
   for (size_t i = 0; i < op_execs_.size(); ++i) {
-    if (op_execs_[i]) op_execs_[i]();
+    if (op_execs_[i]) {
+      PerformOp(i);
+    }
+  }
+}
+
+void GraphExecutor::RematerializeTensor(size_t nid) { PerformOp(nid); }

Review Comment:
   I will try to subclass it.



##########
src/runtime/graph_executor/graph_executor.cc:
##########
@@ -422,14 +527,31 @@ void GraphExecutor::SetupStorage() {
   for (size_t i = 0; i < data_entry_.size(); ++i) {
     int storage_id = attrs_.storage_id[i];
     ICHECK_LT(static_cast<size_t>(storage_id), storage_pool_.size());
-    data_entry_[i] = storage_pool_[storage_id].CreateView(attrs_.shape[i], vtype[i]);
+    // just make view for allocated parameters and tensors in cpu
+    if (pool_entry_[storage_id].device_type == static_cast<int>(kDLCPU)){
+      data_entry_[i] = storage_pool_[storage_id].CreateView(attrs_.shape[i], vtype[i]);
+    }
 
-    const DLTensor* tmp = data_entry_[i].operator->();
-    data_alignment_[i] = details::GetDataAlignment(*tmp);
+    // determine memory usage using grabbed vtype
+    size_t align = (vtype[i].bits / 8) * vtype[i].lanes;
+    data_alignment_[i] = align < kAllocAlignment ? kAllocAlignment : align;
   }
 }
 
 void GraphExecutor::SetupOpExecs() {
+  ref_cnt.resize(num_node_entries(),0);
+  nodes_not_evicted_ = {};
+  for(size_t nid = 0;nid < num_node_entries();nid++){
+    auto slice_node = nodes_[nid];
+    if (slice_node.name.find("slice") != std::string::npos){

Review Comment:
   In my experiment, I found that if I do not leave the slice nodes alone the result of my computation would go wrong, so I encoded this logic to ensure the correctness of my computation.



##########
src/runtime/graph_executor/graph_executor.cc:
##########
@@ -59,10 +61,98 @@ inline size_t GetDataAlignment(const DLTensor& arr) {
 void GraphExecutor::Run() {
   // setup the array and requirements.
   for (size_t i = 0; i < op_execs_.size(); ++i) {
-    if (op_execs_[i]) op_execs_[i]();
+    if (op_execs_[i]) {
+      PerformOp(i);
+    }
+  }
+}
+
+void GraphExecutor::RematerializeTensor(size_t nid) { PerformOp(nid); }
+
+void GraphExecutor::PerformOp(size_t nid) {
+  auto node = nodes_[nid];
+  // first we need to get all the input tensors and add the reference count
+  // if they are not in memory, rematerialize them
+  for (const auto& input : node.inputs) {
+    // also we do not care about cpu
+    if (pool_entry_[get_sid(entry_id(input))].device_type == (int)(kDLCPU)) continue;
+    // we do not need to care about input and parameter
+    if (std::find(input_nodes_.begin(), input_nodes_.end(), entry_id(input)) != input_nodes_.end())
+      continue;
+    if (std::find(nodes_not_evicted_.begin(), nodes_not_evicted_.end(), entry_id(input)) != nodes_not_evicted_.end())
+      continue;
+    // otherwise we should rematerialize the input tensors that are not in the memory
+    // and update the reference count at the same time
+
+    if (is_evicted(entry_id(input))) {
+        RematerializeTensor(entry_id(input));
+    }
+    ref_cnt[entry_id(input)] += 1;
+  }
+  // now all of the inputs are recovered, we can set up to compute the location of input/output
+  // compute current input memory location of the tensor
+  std::unordered_set<uint32_t> input_node_eids;
+  for (unsigned int nid_ : input_nodes_) {
+    input_node_eids.insert(entry_id(nid_, 0));
+  }
+  std::unordered_set<uint32_t> output_node_eids;
+  for (auto& output : outputs_) {
+    output_node_eids.insert(entry_id(output));
+  }
+
+  const auto& inode = nodes_[nid];
+  std::vector<DLTensor> args;
+
+  for (const auto& e : inode.inputs) {
+    uint32_t eid = this->entry_id(e);
+    args.push_back(*(data_entry_[eid].operator->()));
+  }
+
+  /*
+   * Note that some output tensors might not be allocated because some op might produce plural
+   * values Thus we should check if the outputs were allocated and set up the data entry as well
+   */
+  for (uint32_t index = 0; index < inode.param.num_outputs; ++index) {
+    uint32_t eid = this->entry_id(nid, index);
+    set_tensor(eid);
+    args.push_back(*(data_entry_[eid].operator->()));
+  }
+
+  // Now we should update the memory location in TVM Operator
+  std::shared_ptr<OpArgs> updated_args = UpdateTVMOp(nid, nodes_[nid].param, args);
+
+  UpdateInputOutputTensors(input_node_eids, output_node_eids, nid, updated_args);
+
+  // now we can execute the operator
+  op_execs_[nid]();
+  // at last, we should perform eviction
+  PerformEviction(nid);
+}
+
+void GraphExecutor::PerformEviction(size_t nid) {

Review Comment:
   I will change 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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] ruxiliang commented on pull request #11286: Dynamic management

Posted by GitBox <gi...@apache.org>.
ruxiliang commented on PR #11286:
URL: https://github.com/apache/tvm/pull/11286#issuecomment-1132395554

   > exciting to see this work making it's way in! have some initial comments and requests after skimming, will probably revisit this at some later point for further review (but will defer to @MarisaKirisame on the more technical points)
   > 
   > one thing I noticed is that the core eviction & heuristic logic is missing from this PR. not a problem, but are you planning on adding that in a followup?
   
   Yes. Actually I wrote a [simple version](https://github.com/ruxiliang/TVMDynMem/tree/simple-heuristic) but due to time limit I just implemented very simple heurisitc


-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] slyubomirsky commented on a diff in pull request #11286: Dynamic management

Posted by GitBox <gi...@apache.org>.
slyubomirsky commented on code in PR #11286:
URL: https://github.com/apache/tvm/pull/11286#discussion_r881196061


##########
src/runtime/graph_executor/graph_executor.cc:
##########
@@ -448,74 +570,65 @@ void GraphExecutor::SetupOpExecs() {
   for (uint32_t nid = 0; nid < this->GetNumOfNodes(); ++nid) {
     const auto& inode = nodes_[nid];
     if (inode.op_type == "null") continue;
-    std::vector<DLTensor> args;
-    for (const auto& e : inode.inputs) {
-      uint32_t eid = this->entry_id(e);
-      args.push_back(*(data_entry_[eid].operator->()));
-    }
-    for (uint32_t index = 0; index < inode.param.num_outputs; ++index) {
-      uint32_t eid = this->entry_id(nid, index);
-      args.push_back(*(data_entry_[eid].operator->()));
-    }
-    ICHECK(inode.op_type == "tvm_op") << "Can only take tvm_op as op";
+    std::vector<DLTensor> args = {};
 
     std::shared_ptr<OpArgs> op_args = nullptr;
-    std::tie(op_execs_[nid], op_args) = CreateTVMOp(inode.param, args);
-
-    for (size_t i = 0; i < inode.inputs.size(); i++) {
-      uint32_t input_eid = this->entry_id(inode.inputs[i]);
-      // check if op input is model input
-      if (input_node_eids.count(input_eid) > 0) {
-        input_dltensors_[input_eid].push_back(
-            static_cast<DLTensor*>(op_args->arg_values[i].v_handle));
+    std::tie(op_execs_[nid], op_args) = CreateTVMOp(nid, inode.param, args);
+  }
+}
+
+void GraphExecutor::UpdateInputOutputTensors(const std::unordered_set<uint32_t>& input_node_eids,
+                                             const std::unordered_set<uint32_t>& output_node_eids,
+                                             uint32_t nid,
+                                             std::shared_ptr<GraphExecutor::OpArgs> op_args) {
+  const auto& inode = nodes_[nid];
+  for (size_t i = 0; i < inode.inputs.size(); i++) {
+    uint32_t input_eid = this->entry_id(inode.inputs[i]);
+    // check if op input is model input
+    if (input_node_eids.count(input_eid) > 0) {
+      auto& eid_input = input_dltensors_[input_eid];
+      auto input_tensor = static_cast<DLTensor*>(op_args->arg_values[i].v_handle);
+      if (std::find(eid_input.begin(), eid_input.end(), input_tensor) == eid_input.end()) {
+        eid_input.push_back(input_tensor);
       }
-      // check if any model output is the input of the op
-      if (output_node_eids.count(input_eid) > 0) {
-        both_output_opinput_dltensors_[input_eid].push_back(
-            static_cast<DLTensor*>(op_args->arg_values[i].v_handle));
+    }
+    // check if any model output is the input of the op
+    if (output_node_eids.count(input_eid) > 0) {
+      auto& eid_both = both_output_opinput_dltensors_[input_eid];
+      auto both_tensor = static_cast<DLTensor*>(op_args->arg_values[i].v_handle);
+      if(std::find(eid_both.begin(), eid_both.end(), both_tensor) == eid_both.end()){
+        eid_both.push_back(both_tensor);
       }
     }
+  }
 
-    for (uint32_t i = inode.inputs.size(); i < inode.inputs.size() + inode.param.num_outputs; ++i) {
-      uint32_t output_eid = this->entry_id(nid, i - inode.inputs.size());
-      // check if op output is model output
-      if (output_node_eids.count(output_eid) > 0) {
-        output_dltensors_[output_eid].push_back(
-            static_cast<DLTensor*>(op_args->arg_values[i].v_handle));
+  for (uint32_t i = inode.inputs.size(); i < inode.inputs.size() + inode.param.num_outputs; ++i) {
+    uint32_t output_eid = this->entry_id(nid, i - inode.inputs.size());
+    // check if op output is model output
+    if (output_node_eids.count(output_eid) > 0) {
+      auto& eid_out = output_dltensors_[output_eid];
+      auto out_tensor = static_cast<DLTensor*>(op_args->arg_values[i].v_handle);
+      if(std::find(eid_out.begin(), eid_out.end(), out_tensor) == eid_out.end()){
+        eid_out.push_back(out_tensor);
       }
     }
   }
 }
-
 std::pair<std::function<void()>, std::shared_ptr<GraphExecutor::OpArgs> >
-GraphExecutor::CreateTVMOp(const TVMOpParam& param, const std::vector<DLTensor>& args) {
+GraphExecutor::CreateTVMOp(uint32_t nid, const TVMOpParam& param,
+                           const std::vector<DLTensor>& args) {
   std::shared_ptr<GraphExecutor::OpArgs> arg_ptr = std::make_shared<GraphExecutor::OpArgs>();
-  // setup address.
-  arg_ptr->args = args;
-  if (param.flatten_data) {
-    arg_ptr->shape_data.resize(arg_ptr->args.size());
-  }
-  for (size_t i = 0; i < arg_ptr->args.size(); ++i) {
-    TVMValue v;
-    DLTensor* t = &arg_ptr->args[i];
-    v.v_handle = t;
-    arg_ptr->arg_values.push_back(v);
-    arg_ptr->arg_tcodes.push_back(kTVMDLTensorHandle);
-    if (param.flatten_data) {
-      arg_ptr->shape_data[i] =
-          std::accumulate(t->shape, t->shape + t->ndim, 1, std::multiplies<int64_t>());
-      t->ndim = 1;
-      t->shape = &(arg_ptr->shape_data[i]);
-    }
-  }
+  // set up the memory location map
+  op_mem_location_map_[nid] = arg_ptr;
 
   if (param.func_name == "__nop") {
     return {[]() {}, arg_ptr};
   } else if (param.func_name == "__copy") {
     // Perform cross device data copy.
     // Directly copy data from the input to the output.
     // TODO(mbs): device_copy cleanup.
-    auto fexec = [arg_ptr]() {
+    auto fexec = [nid, this]() {
+      auto arg_ptr = this->op_mem_location_map_[nid];

Review Comment:
   Perhaps it might be more economical to define `arg_ptr` outside the lambda and capture only that? (Same with the other lambdas that capture `this`.) I doubt it's a serious issue but it may be worth checking.



-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] MarisaKirisame commented on pull request #11286: Dynamic management

Posted by GitBox <gi...@apache.org>.
MarisaKirisame commented on PR #11286:
URL: https://github.com/apache/tvm/pull/11286#issuecomment-1126681454

   Thx for the PR! Can you add a test case, that dictates how one should use the feature you add?


-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] slyubomirsky commented on pull request #11286: Dynamic management

Posted by GitBox <gi...@apache.org>.
slyubomirsky commented on PR #11286:
URL: https://github.com/apache/tvm/pull/11286#issuecomment-1136702157

   I am also very excited to see steps toward an implementation of DTR in TVM. I agree with Altan's comments, especially on making the dynamically managed executor a subclass of the graph executor (it shouldn't be default behavior). I think it would make sense to include a simple heuristic and some kind of interface for choosing this executor from Python so we could have end-to-end tests.


-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] slyubomirsky commented on a diff in pull request #11286: Dynamic management

Posted by GitBox <gi...@apache.org>.
slyubomirsky commented on code in PR #11286:
URL: https://github.com/apache/tvm/pull/11286#discussion_r881195310


##########
src/runtime/graph_executor/graph_executor.cc:
##########
@@ -422,14 +527,31 @@ void GraphExecutor::SetupStorage() {
   for (size_t i = 0; i < data_entry_.size(); ++i) {
     int storage_id = attrs_.storage_id[i];
     ICHECK_LT(static_cast<size_t>(storage_id), storage_pool_.size());
-    data_entry_[i] = storage_pool_[storage_id].CreateView(attrs_.shape[i], vtype[i]);
+    // just make view for allocated parameters and tensors in cpu
+    if (pool_entry_[storage_id].device_type == static_cast<int>(kDLCPU)){
+      data_entry_[i] = storage_pool_[storage_id].CreateView(attrs_.shape[i], vtype[i]);
+    }
 
-    const DLTensor* tmp = data_entry_[i].operator->();
-    data_alignment_[i] = details::GetDataAlignment(*tmp);
+    // determine memory usage using grabbed vtype
+    size_t align = (vtype[i].bits / 8) * vtype[i].lanes;
+    data_alignment_[i] = align < kAllocAlignment ? kAllocAlignment : align;
   }
 }
 
 void GraphExecutor::SetupOpExecs() {
+  ref_cnt.resize(num_node_entries(),0);
+  nodes_not_evicted_ = {};
+  for(size_t nid = 0;nid < num_node_entries();nid++){
+    auto slice_node = nodes_[nid];
+    if (slice_node.name.find("slice") != std::string::npos){

Review Comment:
   Can you include some of these cases as unit tests? I can't think of a reason for split to be special and it would be good to debug what happened with 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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Dynamic management [tvm]

Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen closed pull request #11286: Dynamic management
URL: https://github.com/apache/tvm/pull/11286


-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org