[Unity][CodeGen] RunCodegen based on externally-exposed functions

Prior to this commit, `relax.transform.RunCodegen` required a list of
entry functions for a module, defaulting to `"main"` if not specified.
The list of entry functions is duplicate information that could be
inferred from the module, and should not be required from the user.
This commit updates `RunCodegen` to treat all externally-exposed
functions as entry points, in the same manner as
`DeadCodeElimination`.

For backwards compatibility, the `entry_functions` argument is still
accepted, and is used to augment the list of externally-exposed
functions.

Author: Lunderberg

python/tvm/relax/transform/transform.py

@@ -574,7 +574,8 @@ def RunCodegen(
         The registered pass to remove unused functions.
     """
     if entry_functions is None:
-        entry_functions = ["main"]
+        entry_functions = []
+
     # enable cutlass byoc registries
     # pylint: disable=unused-import,import-outside-toplevel
     from tvm.contrib import cutlass as _cutlass

src/relax/transform/run_codegen.cc

@@ -28,6 +28,7 @@
 
 #include <iostream>
 
+#include "../../support/ordered_set.h"
 #include "utils.h"
 
 namespace tvm {
@@ -39,12 +40,39 @@ class CodeGenRunner : ExprMutator {
 
   explicit CodeGenRunner(IRModule mod) : ExprMutator(mod) {}
 
-  IRModule Run(Optional<Map<String, OptionMap>> target_options, Array<String> entry_functions) {
+  IRModule Run(Optional<Map<String, OptionMap>> target_options,
+               Array<String> entry_function_names) {
     IRModule mod = builder_->GetContextIRModule();
-    for (const String& entry_func_name : entry_functions) {
-      auto entry_func = mod->Lookup(entry_func_name);
-      auto gvar = mod->GetGlobalVar(entry_func_name);
-      builder_->UpdateFunction(gvar, Downcast<BaseFunc>(VisitExpr(entry_func)));
+
+    support::OrderedSet<GlobalVar> entry_functions;
+    // Any user-provided functions are treated as entry functions.
+    for (const auto& name : entry_function_names) {
+      entry_functions.insert(mod->GetGlobalVar(name));
+    }
+
+    // In addtion, any externally-exposed function that does not
+    // belong to a specific codegen may be an entry function.  These
+    // are added in alphabetical order, to ensure consistent order of
+    // evaluation for debug/test purposes.
+    {
+      std::vector<GlobalVar> attr_entry_functions;
+      for (const auto& [gv, func] : mod->functions) {
+        if (func->GetLinkageType() == LinkageType::kExternal &&
+            !func->GetAttr<String>(attr::kCodegen) && func->IsInstance<relax::FunctionNode>()) {
+          attr_entry_functions.push_back(gv);
+        }
+      }
+      std::sort(attr_entry_functions.begin(), attr_entry_functions.end(),
+                [](const auto& gvar_a, const auto& gvar_b) {
+                  return gvar_a->name_hint > gvar_b->name_hint;
+                });
+      for (const auto& gvar : attr_entry_functions) {
+        entry_functions.insert(gvar);
+      }
+    }
+
+    for (const auto& gvar : entry_functions) {
+      builder_->UpdateFunction(gvar, Downcast<BaseFunc>(VisitExpr(mod->Lookup(gvar))));
     }
 
     auto ext_mods = InvokeCodegen(mod, target_options.value_or({}));
@@ -65,7 +93,7 @@ class CodeGenRunner : ExprMutator {
     }
 
     // TODO(@tvm-team): Implicit pass dependency. Revisit when we have a better way to handle this.
-    return DeadCodeElimination(out_mod, entry_functions);
+    return DeadCodeElimination(out_mod, entry_function_names);
   }
 
   using ExprMutator::VisitExpr_;

tests/python/relax/test_transform_codegen_pass.py

@@ -48,21 +48,21 @@
 dev = tvm.cuda()
 
 
-def check_executable(exec, dev, inputs, expected):
+def check_executable(exec, dev, inputs, expected, entry_func_name):
     vm = relax.VirtualMachine(exec, dev)
-    out = vm["main"](*inputs)
+    out = vm[entry_func_name](*inputs)
     tvm.testing.assert_allclose(out.numpy(), expected.numpy(), atol=1e-5, rtol=1e-5)
 
 
-def check_roundtrip(exec0, dev, inputs, expected):
+def check_roundtrip(exec0, dev, inputs, expected, entry_func_name="main"):
     exec0.mod.export_library("exec.so")
     exec1 = tvm.runtime.load_module("exec.so")
     os.remove("exec.so")
     assert exec0.stats() == exec1["stats"]()
     assert exec0.as_text() == exec1["as_text"]()
 
-    check_executable(exec0, dev, inputs, expected)
-    check_executable(exec1, dev, inputs, expected)
+    check_executable(exec0, dev, inputs, expected, entry_func_name)
+    check_executable(exec1, dev, inputs, expected, entry_func_name)
 
 
 def gen_ground_truth(mod, target, dev, inputs):
@@ -113,10 +113,17 @@ def setup_test():
     return mod, inputs, expected
 
 
+entry_func_name = tvm.testing.parameter("main", "func")
+
+
 @tvm.testing.requires_gpu
-def test_tensorrt_only():
+def test_tensorrt_only(entry_func_name):
     mod, inputs, expected = setup_test()
 
+    if entry_func_name != "main":
+        mod[entry_func_name] = mod
+        del mod["main"]
+
     # Define patterns that we want to offload to byoc
     # This test will offload entire model
     # Thus, define patterns for both `multiply` and `add` ops
@@ -135,7 +142,7 @@ def test_tensorrt_only():
 
     ex0 = relax.build(new_mod, target, params={})
     # Sanity check for the correctness and roundtrip
-    check_roundtrip(ex0, dev, inputs, expected)
+    check_roundtrip(ex0, dev, inputs, expected, entry_func_name)
 
 
 @tvm.testing.requires_gpu
@@ -248,6 +255,28 @@ def test_multiple_calls_same_extern():
     tvm.ir.assert_structural_equal(mod["main"], Conv2dx2_after["main"])
 
 
+def test_default_entry_func():
+    """The entry function is not necessarily named "main"
+
+    Like `test_multiple_calls_same_extern`, but the main function is
+    named "func".
+    """
+    before_with_main = Conv2dx2
+    after_with_main = relax.transform.RunCodegen()(before_with_main)
+
+    def rename_main(mod):
+        mod = mod.clone()
+        mod["func"] = mod["main"].with_attr("global_symbol", "func")
+        del mod["main"]
+        return mod
+
+    before_with_func = rename_main(before_with_main)
+    expected_with_func = rename_main(after_with_main)
+    after_with_func = relax.transform.RunCodegen()(before_with_func)
+
+    tvm.ir.assert_structural_equal(expected_with_func["func"], after_with_func["func"])
+
+
 def test_dynamic_shape():
     import tvm.relax.backend.contrib.cublas