(tvm) branch main updated: [TIR] ConvertSSA process entry func first (#16236)

[tq...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new b0e146f767 [TIR] ConvertSSA process entry func first (#16236)
b0e146f767 is described below

commit b0e146f7676ea35d26f22c81e53f6dfc312b4514
Author: Hongyi Jin <ji...@gmail.com>
AuthorDate: Thu Dec 14 08:07:20 2023 -0500

    [TIR] ConvertSSA process entry func first (#16236)
    
    * convert ssa process entry func first
    
    * add doc
---
 src/tir/transforms/ir_utils.cc | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/src/tir/transforms/ir_utils.cc b/src/tir/transforms/ir_utils.cc
index 25c10dd682..6b681c07e5 100644
--- a/src/tir/transforms/ir_utils.cc
+++ b/src/tir/transforms/ir_utils.cc
@@ -715,8 +715,13 @@ Pass ConvertSSA() {
     tir::IRConvertSSA converter;
     Map<GlobalVar, BaseFunc> functions;
     bool made_change = false;
+    // FIXME: This is just a temporal workaround to ensure free vars
+    // in device function have the same pointer as the host function
     for (auto [gvar, base_func] : mod->functions) {
       if (auto* ptr = base_func.as<tir::PrimFuncNode>()) {
+        if (!ptr->HasNonzeroAttr(tir::attr::kIsEntryFunc)) {
+          continue;
+        }
         auto updated = converter.VisitPrimFunc(GetRef<tir::PrimFunc>(ptr));
         if (!updated.same_as(base_func)) {
           made_change = true;
@@ -725,6 +730,19 @@ Pass ConvertSSA() {
       }
       functions.Set(gvar, base_func);
     }
+    for (auto [gvar, base_func] : mod->functions) {
+      if (auto* ptr = base_func.as<tir::PrimFuncNode>()) {
+        if (ptr->HasNonzeroAttr(tir::attr::kIsEntryFunc)) {
+          continue;
+        }
+        auto updated = converter.VisitPrimFunc(GetRef<tir::PrimFunc>(ptr));
+        if (!updated.same_as(base_func)) {
+          made_change = true;
+          base_func = updated;
+        }
+        functions.Set(gvar, base_func);
+      }
+    }
     if (made_change) {
       mod.CopyOnWrite()->functions = std::move(functions);
     }