You are viewing a plain text version of this content. The canonical link for it is here.

Posted to commits@tvm.apache.org by "Lunderberg (via GitHub)" <gi...@apache.org> on 2024/01/25 15:14:45 UTC

[PR] [Transform] Preserve all reachable functions in FuseTIR [tvm]

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

   Prior to this commit, the `FuseTIR` pass had custom logic to determine whether a `PrimFunc` should be kept in the output `IRModule`.  This commit replaces this check in `FuseTIR` with a post-processing by `DeadCodeElimination`.


-- 
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] [Transform] Preserve all reachable functions in FuseTIR [tvm]

Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.

slyubomirsky commented on code in PR #16471:
URL: https://github.com/apache/tvm/pull/16471#discussion_r1467148384


##########
src/relax/transform/fuse_tir.cc:
##########
@@ -950,63 +940,66 @@ class TIRFuseMutator : public ExprMutator {
 
     if (call->op->IsInstance<GlobalVarNode>()) {
       // Case 1. It is a relax cross function call
-      GlobalVar old_gv = Downcast<GlobalVar>(call->op);
-      auto relax_func = Downcast<Function>(mod_->Lookup(old_gv));
-      auto it = fused_tir_funcs_.find(old_gv);
-      if (it != fused_tir_funcs_.end()) {
-        const tir::PrimFunc& fused_tir = (*it).second;
-        // Case 1.1. It calls a primitive relax function, update the call into a call_tir
-        GlobalVar fused_tir_gv = this->builder_->AddFunction(fused_tir, old_gv->name_hint);
-        // Step a. Flatten all args since call_tir does not support Tuple value.
-        Array<Expr> arg_list;
-        Array<PrimExpr> tir_vars;
-        for (size_t i = 0; i < call->args.size(); ++i) {
-          auto arg = call->args[i];
-          auto sinfo = GetStructInfo(arg);
-
-          ICHECK(!relax_func->params[i]->struct_info_->IsInstance<TupleStructInfoNode>() &&
-                 !sinfo.as<TupleStructInfoNode>())
-              << "InternalError: "
-              << "All tuple parameters should be expanded before this point in FuseTIR.  "
-              << "However, argument " << arg << " with struct info " << arg->struct_info_
-              << " is passed as argument " << i << " to Primitive Relax function " << old_gv
-              << ", which expects parameter " << relax_func->params[i] << " to have struct info "
-              << relax_func->params[i]->struct_info_;
-
-          if (const auto* shape = sinfo.as<ShapeStructInfoNode>()) {
-            CHECK(shape->values.defined())
-                << "FuseTIR requires all shape input has struct_info value.";
-            for (const PrimExpr& prim_value : shape->values.value()) {
-              CHECK(prim_value->IsInstance<tir::VarNode>())
-                  << "All shape inputs are expected to be single tir var.";
-              tir_vars.push_back(prim_value);
-            }
-          } else if (const auto* prim_value = sinfo.as<PrimStructInfoNode>()) {
-            CHECK(prim_value->value.defined())
-                << "FuseTIR requires all R.Prim arguments to have a known value.";
-            PrimExpr expr = prim_value->value.value();
-            CHECK(expr->IsInstance<tir::VarNode>())
-                << "FuseTIR currently requires all R.Prim arguments to provide a single tir::Var.";
-            tir_vars.push_back(expr);
-
-          } else {
-            arg_list.push_back(arg);
+      GlobalVar old_gvar = Downcast<GlobalVar>(call->op);
+      auto relax_func = Downcast<Function>(orig_mod_->Lookup(old_gvar));
+
+      if (!relax_func->HasNonzeroAttr(attr::kPrimitive)) {
+        // The callee is not a primitive function, no need to fuse.
+        return call;
+      }
+
+      tir::PrimFunc fused_tir = FusedTIRConstructor::GetFusedTIR(orig_mod_, old_gvar);

Review Comment:
   Are you sure we aren't likely to call this more than once per global var? As long as it's not likely or costly, this approach is probably fine.



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