[GitHub] [tvm] ashutosh-arm commented on a diff in pull request #11091: [AOT] Enable A-Normal Form in the AOT executor

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

ashutosh-arm commented on code in PR #11091:
URL: https://github.com/apache/tvm/pull/11091#discussion_r858867879


##########
src/relay/backend/contrib/cmsisnn/relay_to_tir.cc:
##########
@@ -655,19 +655,61 @@ class RelayToTIRVisitor : public MixedModeMutator {
     return Call(new_global_var, call->args, call->attrs, call->type_args, call->span);
   }
 
-  Expr Rewrite_(const CallNode* pre, const Expr& post) override {
-    if (const CallNode* call = post.as<CallNode>()) {
-      auto* func = call->op.as<FunctionNode>();
-      if (func == nullptr) {
-        return post;
+  Expr VisitExpr_(const LetNode* op) final {
+    auto pre_visit = [this](const LetNode* op) {
+      Expr var = this->VisitExpr(op->var);
+      Expr value = this->VisitExpr(op->value);
+      // outlineable function no longer needs let binding
+      if (this->CanOutlineExpr(value)) {
+        this->memo_[var] = value;
+      }
+    };
+    auto post_visit = [this](const LetNode* op) {
+      // Rely on the Memoizer to cache pre-visit values
+      Expr value = this->VisitExpr(op->value);
+      Expr body = this->VisitExpr(op->body);
+      auto expr = GetRef<Expr>(op);
+      // drop the let binding
+      if (this->CanOutlineExpr(value)) {
+        this->memo_[expr] = this->VisitExpr(op->body);
+      } else {
+        Var var = Downcast<Var>(this->VisitExpr(op->var));
+        if (var.same_as(op->var) && value.same_as(op->value) && body.same_as(op->body)) {
+          this->memo_[expr] = expr;
+        } else {
+          this->memo_[expr] = Let(var, value, body);
+        }
       }
+    };
+    ExpandANormalForm(op, pre_visit, post_visit);
+    return memo_[GetRef<Expr>(op)];
+  }
 
-      auto codegen_name = func->GetAttr<String>(attr::kCompiler);
-      if (codegen_name.defined() && codegen_name == "cmsis-nn") {
-        const CallNode* inner_call = func->body.as<CallNode>();
+  bool CanOutlineExpr(const Expr& expr) {
+    // TODO(@lhutton1): This behaviour is similar to the OutlineCompilerFunctions pass
+    // we could reuse this functionality by separating outlining and lowering in this
+    // pass.

Review Comment:
   I remember these functions were outlined before we moved to target hooks. Some optimizations during before they are lowered must have made them inlined. Just for my understanding, what is the advantage of separating them?



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