areusch commented on a change in pull request #8096:
URL: https://github.com/apache/tvm/pull/8096#discussion_r649381183
##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -44,52 +45,184 @@ namespace tvm {
namespace relay {
namespace backend {
+/**
+ * Struct to contain information about the intermediate tensors in the
+ * runner function
+ */
+struct StorageInfo {
+ /*! \brief storage integer identifier of the particular intermediate buffer
*/
+ int sid;
+ /*! \brief exact size of the temporary */
+ int size_bytes;
+ /*! \brief device type of the intermediate tensor */
+ int dev_type;
+};
+
using IntegerArray = Array<Integer>;
using TargetsMap = std::unordered_map<int, Target>;
+using StorageMap = std::unordered_map<Expr, std::vector<StorageInfo>,
runtime::ObjectPtrHash,
+ runtime::ObjectPtrEqual>;
-class AotReturnSidVisitor : public ExprVisitor {
+/**
+ * This is an on demand allocator for AOT. A new temporary
+ * (storage allocator identifier) is allocated for each operation.
+ */
+class AOTOnDemandAllocator : public ExprVisitor {
Review comment:
i guess, do we want to add any asserts here to ensure that we don't see
top-level nodes other than Constant, CallNode, VarNode, or TupleGetItemNode?
what if someone modifies the AOT TIR codegen to add some new pattern but
forgets to update this? it seems like something should fail then.
##########
File path: tests/python/relay/aot/test_crt_aot.py
##########
@@ -364,5 +364,27 @@ def test_byoc_utvm(use_calculated_workspaces):
compile_and_run(mod, input_list, output_list, use_calculated_workspaces)
+def test_quant_mobilenet_tfl():
Review comment:
ok so here's my concern with this--it's great to test on real-world test
cases, but we are really only certain that this test case exercises the logic
path we want it to because, at the time of this PR, we manually traced it
through and verified that this model triggers the particular scenario we care
about. most ideally there would be a contrived unit test, but it can be tricky
to write given the number of moving pieces. so i'm okay with using mobilenet as
an end-to-end thing, but i'd like you to add some asserts to this test case
that verifies that the logic path you're trying to test (e.g. avoiding reuse of
the output buffer) is actually triggered, such that if someone modifies
AOTExecutorCodegen in the future, this test would fail if that is no longer
true.
##########
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:
i guess another way to view this though is that this PR is adding
reference tracking, but reference tracking only works with one particular TIR
pattern. i agree with you that adding full reference tracking is a lot to ask
particularly considering the goals of this PR. However, there's no
documentation (by which I mean no unit test) as to what specifically needs to
be tracked. can you add a unit test for the changes you've made to
StorageRewrite?
--
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:
[email protected]
