[GitHub] [tvm] giuseros commented on a change in pull request #8096: Decoupling AOT from graph memory planner

[GitBox <gi...@apache.org>]

giuseros commented on a change in pull request #8096:
URL: https://github.com/apache/tvm/pull/8096#discussion_r646944601



##########
File path: src/tir/transforms/storage_rewrite.cc
##########
@@ -138,6 +138,34 @@ class LinearAccessPatternFinder final : public StmtExprVisitor {
     if (op->op.same_as(builtin::address_of())) {
       const LoadNode* l = op->args[0].as<LoadNode>();
       this->VisitExpr(l->index);
+    } else if (op->op.same_as(builtin::tvm_call_cpacked())) {
+      // Recall that the arguments of a tvm_call_cpacked are passed as
+      // TVMValues. But a TVMValue is only a container, that points to
+      // a real buffer previously allocated. We need to signal that those
+      // buffers need to be live at the same time (i.e., cannot be overridden)
+      Array<PrimExpr> args = op->args;
+      for (auto arg : args) {
+        const VarNode* var = arg.as<VarNode>();
+        if (value_to_alloc_.find(var) != value_to_alloc_.end()) {
+          auto allocs = value_to_alloc_[var];
+          for (const VarNode* alloc : allocs) {
+            VisitExpr_(alloc);
+          }
+        } else {
+          this->VisitExpr(arg);
+        }
+      }
+    } else if (op->op.same_as(builtin::tvm_struct_set())) {

Review comment:
       So, for called_packed, we don't care about accessing the data (`tvm_struct_get`), but only setting the data of the TVMValue. But  in general it is true that other patterns might be problematic if used. 
   
   I am not sure about `tvm_struct_get` in particular, because you would use the (real) data you extracted, while with `tvm_struct_set` you  use the real data to set the structure and then you pass the structure in the call. This change is only to annotate that the struct is a mere container, and that the rewriter needs to focus on the real data. 
   
   I would say that this PR is not trying to fix the rewriter on all the possible corner cases. This particular corner case is affecting AOT, and it needs to be fixed. If there are other corner cases, I would fix them in a separate PR. 




-- 
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.

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