You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2025/08/12 14:28:06 UTC
(tvm) branch main-mod created (now 4f5f2d9571)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a change to branch main-mod
in repository https://gitbox.apache.org/repos/asf/tvm.git
at 4f5f2d9571 [REFATOR] Phase out entry_func
This branch includes the following new commits:
new 4f5f2d9571 [REFATOR] Phase out entry_func
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(tvm) 01/01: [REFATOR] Phase out entry_func
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main-mod
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 4f5f2d9571f2ec8dde7ee5570752c20cafbcc256
Author: tqchen <ti...@gmail.com>
AuthorDate: Tue Aug 12 10:22:42 2025 -0400
[REFATOR] Phase out entry_func
This PR phases out the entry function. Previously module itself have a default entry function.
Supporting such feature requires extra duplicated code and indirection.
Most of our use cases can move toward explicitly naming the function as mod["main"],
so this PR phases out this behavior.
---
include/tvm/runtime/module.h | 2 -
include/tvm/tir/function.h | 10 -----
include/tvm/tir/transform.h | 6 ---
jvm/core/src/main/java/org/apache/tvm/Module.java | 2 +-
python/tvm/runtime/executable.py | 4 --
python/tvm/runtime/module.py | 22 -----------
python/tvm/tir/build.py | 3 +-
python/tvm/tir/pipeline.py | 1 -
python/tvm/tir/transform/transform.py | 11 ------
src/meta_schedule/arg_info.cc | 10 ++---
src/runtime/cuda/cuda_module.cc | 2 +-
src/runtime/library_module.cc | 6 +--
src/runtime/metal/metal_module.mm | 2 +-
src/runtime/opencl/opencl_module.cc | 2 +-
src/runtime/rocm/rocm_module.cc | 2 +-
src/runtime/vulkan/vulkan_wrapped_func.cc | 2 +-
src/target/llvm/codegen_cpu.cc | 4 +-
src/target/llvm/codegen_hexagon.cc | 11 ------
src/target/llvm/llvm_module.cc | 22 ++---------
src/target/source/codegen_c_host.cc | 14 -------
src/tir/ir/transform.cc | 1 -
src/tir/transforms/primfunc_utils.cc | 43 +---------------------
tests/python/codegen/test_target_codegen_device.py | 2 +-
.../test_hexagon/test_async_dma_pipeline.py | 10 +++--
.../contrib/test_hexagon/test_parallel_hvx.py | 2 +-
.../test_hexagon/test_parallel_hvx_load_vtcm.py | 8 +---
.../contrib/test_hexagon/test_parallel_scalar.py | 4 +-
.../contrib/test_hexagon/test_vtcm_bandwidth.py | 8 +++-
.../test_runtime_builtin_kv_cache_transfer.py | 2 +-
...runtime_builtin_paged_attention_kv_cache_cpu.py | 2 +-
..._builtin_paged_attention_kv_cache_flashinfer.py | 2 +-
...ltin_paged_attention_kv_cache_mla_flashinfer.py | 2 +-
...ime_builtin_paged_attention_kv_cache_mla_tir.py | 2 +-
...runtime_builtin_paged_attention_kv_cache_tir.py | 2 +-
.../python/relax/test_runtime_builtin_rnn_state.py | 2 +-
.../tir-transform/test_tir_transform_helpers.py | 31 ----------------
tests/python/tvmscript/test_tvmscript_roundtrip.py | 2 -
37 files changed, 44 insertions(+), 219 deletions(-)
diff --git a/include/tvm/runtime/module.h b/include/tvm/runtime/module.h
index efbaa6508a..05b57de39d 100644
--- a/include/tvm/runtime/module.h
+++ b/include/tvm/runtime/module.h
@@ -296,8 +296,6 @@ constexpr const char* tvm_set_device = "__tvm_set_device";
constexpr const char* tvm_global_barrier_state = "__tvm_global_barrier_state";
/*! \brief Prepare the global barrier before kernels that uses global barrier. */
constexpr const char* tvm_prepare_global_barrier = "__tvm_prepare_global_barrier";
-/*! \brief Placeholder for the module's entry function. */
-constexpr const char* tvm_module_main = "__tvm_main__";
} // namespace symbol
// implementations of inline functions.
diff --git a/include/tvm/tir/function.h b/include/tvm/tir/function.h
index 6ea50e9ae0..ff9a6ff927 100644
--- a/include/tvm/tir/function.h
+++ b/include/tvm/tir/function.h
@@ -316,16 +316,6 @@ constexpr const char* kKernelLaunchParams = "tir.kernel_launch_params";
*/
constexpr const char* kNoAlias = "tir.noalias";
-/*!
- * \brief Mark the function as the entry function of
- * the final generated runtime module.
- *
- * Type: Integer
- *
- * \note There can only be one entry function per module.
- */
-constexpr const char* kIsEntryFunc = "tir.is_entry_func";
-
/*!
* \brief Mark the function as the global function called from the host.
*
diff --git a/include/tvm/tir/transform.h b/include/tvm/tir/transform.h
index eb64d87f95..c7af05e7f2 100644
--- a/include/tvm/tir/transform.h
+++ b/include/tvm/tir/transform.h
@@ -703,12 +703,6 @@ TVM_DLL Pass RenormalizeSplitPattern();
*/
TVM_DLL Pass BindTarget(Target target);
-/*!
- * \brief Set a PrimFunc as the entry point if it is only function in IRModule.
- * \return The pass.
- */
-TVM_DLL Pass AnnotateEntryFunc();
-
/*!
* \brief Filter PrimFuncs with a given condition.
* \return The pass.
diff --git a/jvm/core/src/main/java/org/apache/tvm/Module.java b/jvm/core/src/main/java/org/apache/tvm/Module.java
index 5e78e26ae7..9fa65054f9 100644
--- a/jvm/core/src/main/java/org/apache/tvm/Module.java
+++ b/jvm/core/src/main/java/org/apache/tvm/Module.java
@@ -46,7 +46,7 @@ public class Module extends TVMObject {
}
private Function entry = null;
- private final String entryName = "__tvm_main__";
+ private final String entryName = "__tvm_ffi_main__";
/**
diff --git a/python/tvm/runtime/executable.py b/python/tvm/runtime/executable.py
index b6e13a65a9..a1a6606765 100644
--- a/python/tvm/runtime/executable.py
+++ b/python/tvm/runtime/executable.py
@@ -36,10 +36,6 @@ class Executable:
"""Get the PackedFunc from the jitted module."""
return self.jit().get_function(name, query_imports=True)
- def __call__(self, *args, **kwargs) -> Any:
- """Call the executable."""
- return self.jit().entry_func(*args, **kwargs)
-
def jit(
self,
*,
diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py
index 3dd4de5da0..30f83474dc 100644
--- a/python/tvm/runtime/module.py
+++ b/python/tvm/runtime/module.py
@@ -103,24 +103,8 @@ class Module(tvm.ffi.Object):
def __new__(cls):
instance = super(Module, cls).__new__(cls) # pylint: disable=no-value-for-parameter
- instance.entry_name = "__tvm_main__"
- instance._entry = None
return instance
- @property
- def entry_func(self):
- """Get the entry function
-
- Returns
- -------
- f : tvm.runtime.PackedFunc
- The entry function if exist
- """
- if self._entry:
- return self._entry
- self._entry = self.get_function("__tvm_main__")
- return self._entry
-
def implements_function(self, name, query_imports=False):
"""Returns True if the module has a definition for the global function with name. Note
that has_function(name) does not imply get_function(name) is non-null since the module
@@ -179,12 +163,6 @@ class Module(tvm.ffi.Object):
raise ValueError("Can only take string as function name")
return self.get_function(name)
- def __call__(self, *args):
- if self._entry:
- return self._entry(*args)
- # pylint: disable=not-callable
- return self.entry_func(*args)
-
@property
def type_key(self):
"""Get type key of the module."""
diff --git a/python/tvm/tir/build.py b/python/tvm/tir/build.py
index 98e549cc9c..431c601e72 100644
--- a/python/tvm/tir/build.py
+++ b/python/tvm/tir/build.py
@@ -80,8 +80,7 @@ def split_host_device_mods(mod: IRModule) -> Tuple[IRModule, Dict[Target, IRModu
@T.prim_func
def main(self_handle: T.handle, args: T.handle, num_args: T.int32, result: T.handle):
T.func_attr({"target": T.target({"keys": ["cpu"], "kind": "c"}),
- "calling_conv": 1, # kCPackedFunc for entry functions
- "tir.is_entry_func": True})
+ "calling_conv": 1})
# ... main function implementation
The function will return:
diff --git a/python/tvm/tir/pipeline.py b/python/tvm/tir/pipeline.py
index ae78b05738..1082cd8fac 100644
--- a/python/tvm/tir/pipeline.py
+++ b/python/tvm/tir/pipeline.py
@@ -89,7 +89,6 @@ def default_tir_pipeline():
tir.transform.VerifyVTCMLimit(),
tir.transform.LowerVtcmAlloc(),
tir.transform.VerifyMemory(),
- tir.transform.AnnotateEntryFunc(),
]
)
if bool(config.get("tir.detect_global_barrier", False)):
diff --git a/python/tvm/tir/transform/transform.py b/python/tvm/tir/transform/transform.py
index 93a182ca3b..178a203ca5 100644
--- a/python/tvm/tir/transform/transform.py
+++ b/python/tvm/tir/transform/transform.py
@@ -1018,17 +1018,6 @@ def BindTarget(target):
return _ffi_api.BindTarget(target) # type: ignore
-def AnnotateEntryFunc():
- """Set a PrimFunc as the entry point if it is only function in IRModule.
-
- Returns
- -------
- fpass : tvm.transform.Pass
- The result pass
- """
- return _ffi_api.AnnotateEntryFunc() # type: ignore
-
-
def Filter(fcond: Callable):
"""Filter out PrimFuncs that does not satisfy the given condition.
`fcond` should be a function that takes a primfunc and returns boolean.
diff --git a/src/meta_schedule/arg_info.cc b/src/meta_schedule/arg_info.cc
index 9c2ba084ad..c46a0bf280 100644
--- a/src/meta_schedule/arg_info.cc
+++ b/src/meta_schedule/arg_info.cc
@@ -25,12 +25,12 @@ namespace meta_schedule {
/*!
* \brief Find the entry function of the given IRModule, i.e, functions marked by
- * `tir::attr::kIsEntryFunc`, whose name is `main` or being the only PrimeFunc.
+ * whose name is `main` or being the only PrimeFunc.
* \param mod The IRModule to find the entry function.
* \return The entry function.
*/
inline tir::PrimFunc FindEntryFunc(const IRModule& mod) {
- // Priority 1: PrimFunc marked as `tir::attr::kIsEntryFunc`
+ // Priority 1: PrimFunc marked as `main`
int num_prim_func = 0;
const tir::PrimFuncNode* main_func = nullptr;
const tir::PrimFuncNode* last_func = nullptr;
@@ -39,9 +39,6 @@ inline tir::PrimFunc FindEntryFunc(const IRModule& mod) {
BaseFunc base_func = kv.second;
if (const auto* func = base_func.as<tir::PrimFuncNode>()) {
last_func = func;
- if (func->HasNonzeroAttr(tir::attr::kIsEntryFunc)) {
- return GetRef<tir::PrimFunc>(func);
- }
if (gv->name_hint == "main") {
main_func = func;
}
@@ -57,8 +54,7 @@ inline tir::PrimFunc FindEntryFunc(const IRModule& mod) {
LOG(FATAL) << "ValueError: Cannot find any PrimFunc in the given IRModule: " << mod;
}
if (num_prim_func > 1) {
- LOG(FATAL) << "ValueError: Multiple PrimFuncs exist in the IRModule, but none of them are "
- "annotated with `kIsEntryFunc`, i.e. `tir.is_entry_func`"
+ LOG(FATAL) << "ValueError: Multiple PrimFuncs exist in the IRModule, but none of them are main"
<< mod;
}
return GetRef<tir::PrimFunc>(last_func);
diff --git a/src/runtime/cuda/cuda_module.cc b/src/runtime/cuda/cuda_module.cc
index 2435cccf0a..8d1ed2b2d1 100644
--- a/src/runtime/cuda/cuda_module.cc
+++ b/src/runtime/cuda/cuda_module.cc
@@ -258,7 +258,7 @@ class CUDAPrepGlobalBarrier {
ffi::Function CUDAModuleNode::GetFunction(const String& name,
const ObjectPtr<Object>& sptr_to_self) {
ICHECK_EQ(sptr_to_self.get(), this);
- ICHECK_NE(name, symbol::tvm_module_main) << "Device function do not have main";
+ ICHECK_NE(name, symbol::tvm_ffi_main) << "Device function do not have main";
if (name == symbol::tvm_prepare_global_barrier) {
return ffi::Function(CUDAPrepGlobalBarrier(this, sptr_to_self));
}
diff --git a/src/runtime/library_module.cc b/src/runtime/library_module.cc
index fffac4adea..77a1072c33 100644
--- a/src/runtime/library_module.cc
+++ b/src/runtime/library_module.cc
@@ -50,11 +50,11 @@ class LibraryModuleNode final : public ModuleNode {
ffi::Function GetFunction(const String& name, const ObjectPtr<Object>& sptr_to_self) final {
TVMFFISafeCallType faddr;
- if (name == runtime::symbol::tvm_module_main) {
+ if (name == runtime::symbol::tvm_ffi_main) {
const char* entry_name =
- reinterpret_cast<const char*>(lib_->GetSymbol(runtime::symbol::tvm_module_main));
+ reinterpret_cast<const char*>(lib_->GetSymbol(runtime::symbol::tvm_ffi_main));
ICHECK(entry_name != nullptr)
- << "Symbol " << runtime::symbol::tvm_module_main << " is not presented";
+ << "Symbol " << runtime::symbol::tvm_ffi_main << " is not presented";
faddr = reinterpret_cast<TVMFFISafeCallType>(lib_->GetSymbol(entry_name));
} else {
faddr = reinterpret_cast<TVMFFISafeCallType>(lib_->GetSymbol(name.c_str()));
diff --git a/src/runtime/metal/metal_module.mm b/src/runtime/metal/metal_module.mm
index be36e6197f..a054378e8f 100644
--- a/src/runtime/metal/metal_module.mm
+++ b/src/runtime/metal/metal_module.mm
@@ -264,7 +264,7 @@ ffi::Function MetalModuleNode::GetFunction(const String& name,
ffi::Function ret;
AUTORELEASEPOOL {
ICHECK_EQ(sptr_to_self.get(), this);
- ICHECK_NE(name, symbol::tvm_module_main) << "Device function do not have main";
+ ICHECK_NE(name, symbol::tvm_ffi_main) << "Device function do not have main";
auto it = fmap_.find(name);
if (it == fmap_.end()) {
ret = ffi::Function();
diff --git a/src/runtime/opencl/opencl_module.cc b/src/runtime/opencl/opencl_module.cc
index 19b426d4b4..57621b8609 100644
--- a/src/runtime/opencl/opencl_module.cc
+++ b/src/runtime/opencl/opencl_module.cc
@@ -138,7 +138,7 @@ cl::OpenCLWorkspace* OpenCLModuleNodeBase::GetGlobalWorkspace() {
ffi::Function OpenCLModuleNodeBase::GetFunction(const String& name,
const ObjectPtr<Object>& sptr_to_self) {
ICHECK_EQ(sptr_to_self.get(), this);
- ICHECK_NE(name, symbol::tvm_module_main) << "Device function do not have main";
+ ICHECK_NE(name, symbol::tvm_ffi_main) << "Device function do not have main";
auto it = fmap_.find(name);
if (it == fmap_.end()) return ffi::Function();
const FunctionInfo& info = it->second;
diff --git a/src/runtime/rocm/rocm_module.cc b/src/runtime/rocm/rocm_module.cc
index 791e4b1569..8a71bcbbf7 100644
--- a/src/runtime/rocm/rocm_module.cc
+++ b/src/runtime/rocm/rocm_module.cc
@@ -195,7 +195,7 @@ class ROCMWrappedFunc {
ffi::Function ROCMModuleNode::GetFunction(const String& name,
const ObjectPtr<Object>& sptr_to_self) {
ICHECK_EQ(sptr_to_self.get(), this);
- ICHECK_NE(name, symbol::tvm_module_main) << "Device function do not have main";
+ ICHECK_NE(name, symbol::tvm_ffi_main) << "Device function do not have main";
auto it = fmap_.find(name);
if (it == fmap_.end()) return ffi::Function();
const FunctionInfo& info = it->second;
diff --git a/src/runtime/vulkan/vulkan_wrapped_func.cc b/src/runtime/vulkan/vulkan_wrapped_func.cc
index f4922a1bf0..d863b02cdf 100644
--- a/src/runtime/vulkan/vulkan_wrapped_func.cc
+++ b/src/runtime/vulkan/vulkan_wrapped_func.cc
@@ -208,7 +208,7 @@ VulkanModuleNode::~VulkanModuleNode() {
ffi::Function VulkanModuleNode::GetFunction(const String& name,
const ObjectPtr<Object>& sptr_to_self) {
ICHECK_EQ(sptr_to_self.get(), this);
- ICHECK_NE(name, symbol::tvm_module_main) << "Device function do not have main";
+ ICHECK_NE(name, symbol::tvm_ffi_main) << "Device function do not have main";
auto it = fmap_.find(name);
if (it == fmap_.end()) return ffi::Function();
const FunctionInfo& info = it->second;
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index 4dd24026c0..a04273094c 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -234,7 +234,7 @@ void CodeGenCPU::AddMainFunction(const std::string& entry_func_name) {
llvm::Type* type = llvm::ArrayType::get(t_char_, entry_func_name.length() + 1);
llvm::GlobalVariable* global =
new llvm::GlobalVariable(*module_, type, true, llvm::GlobalValue::WeakAnyLinkage, nullptr,
- runtime::symbol::tvm_module_main);
+ runtime::symbol::tvm_ffi_main);
#if TVM_LLVM_VERSION >= 100
global->setAlignment(llvm::Align(1));
#else
@@ -243,7 +243,7 @@ void CodeGenCPU::AddMainFunction(const std::string& entry_func_name) {
// comdat is needed for windows select any linking to work
// set comdat to Any(weak linking)
if (llvm_target_->GetOrCreateTargetMachine()->getTargetTriple().isOSWindows()) {
- llvm::Comdat* comdat = module_->getOrInsertComdat(runtime::symbol::tvm_module_main);
+ llvm::Comdat* comdat = module_->getOrInsertComdat(runtime::symbol::tvm_ffi_main);
comdat->setSelectionKind(llvm::Comdat::Any);
global->setComdat(comdat);
}
diff --git a/src/target/llvm/codegen_hexagon.cc b/src/target/llvm/codegen_hexagon.cc
index 6f90da3d8a..3d8ed08eee 100644
--- a/src/target/llvm/codegen_hexagon.cc
+++ b/src/target/llvm/codegen_hexagon.cc
@@ -483,9 +483,6 @@ runtime::Module BuildHexagon(IRModule mod, Target target) {
(void)CallOnce;
auto cg = std::make_unique<CodeGenHexagon>();
-
- std::string entry_func;
-
for (auto kv : mod->functions) {
if (!kv.second->IsInstance<PrimFuncNode>()) {
// (@jroesch): we relax constraints here, relax functions will just be ignored.
@@ -493,18 +490,10 @@ runtime::Module BuildHexagon(IRModule mod, Target target) {
continue;
}
auto f = Downcast<PrimFunc>(kv.second);
- if (f->HasNonzeroAttr(tir::attr::kIsEntryFunc)) {
- auto global_symbol = f->GetAttr<String>(tvm::attr::kGlobalSymbol);
- ICHECK(global_symbol.has_value());
- entry_func = global_symbol.value();
- }
}
cg->Init("TVMHexagonModule", llvm_target.get(), std::nullopt, false, false);
cg->AddFunctionsOrdered(mod->functions.begin(), mod->functions.end());
- if (entry_func.length() != 0) {
- cg->AddMainFunction(entry_func);
- }
// Uncomment to get the LLVM module right out of codegen, before optimizations.
// std::cerr << "HexagonModule.0 {\n" << *cg->GetModulePtr() << "}\n";
diff --git a/src/target/llvm/llvm_module.cc b/src/target/llvm/llvm_module.cc
index a9e09652ee..45ede6efef 100644
--- a/src/target/llvm/llvm_module.cc
+++ b/src/target/llvm/llvm_module.cc
@@ -190,15 +190,7 @@ ffi::Function LLVMModuleNode::GetFunction(const String& name,
TVMFFISafeCallType faddr;
With<LLVMTarget> llvm_target(*llvm_instance_, LLVMTarget::GetTargetMetadata(*module_));
- if (name == runtime::symbol::tvm_module_main) {
- const char* entry_name = reinterpret_cast<const char*>(
- GetGlobalAddr(runtime::symbol::tvm_module_main, *llvm_target));
- ICHECK(entry_name != nullptr) << "Symbol " << runtime::symbol::tvm_module_main
- << " is not presented";
- faddr = reinterpret_cast<TVMFFISafeCallType>(GetFunctionAddr(entry_name, *llvm_target));
- } else {
- faddr = reinterpret_cast<TVMFFISafeCallType>(GetFunctionAddr(name, *llvm_target));
- }
+ faddr = reinterpret_cast<TVMFFISafeCallType>(GetFunctionAddr(name, *llvm_target));
if (faddr == nullptr) return ffi::Function();
return tvm::runtime::WrapFFIFunction(faddr, sptr_to_self);
}
@@ -337,15 +329,9 @@ void LLVMModuleNode::Init(const IRModule& mod, const Target& target) {
}
auto f = Downcast<PrimFunc>(kv.second);
auto global_symbol = f->GetAttr<String>(tvm::attr::kGlobalSymbol);
- bool is_entry_func = f->HasNonzeroAttr(tir::attr::kIsEntryFunc);
-
- ICHECK(global_symbol || !is_entry_func) << "The entry func must be exposed externally.";
if (global_symbol) {
function_names_.push_back(global_symbol.value());
- if (is_entry_func) {
- entry_func = global_symbol.value();
- }
}
}
// TODO(@jroesch): follow up on this condition.
@@ -355,11 +341,9 @@ void LLVMModuleNode::Init(const IRModule& mod, const Target& target) {
cg->Init("TVMMod", llvm_target.get(), system_lib_prefix, system_lib_prefix.has_value(), false);
cg->SetFastMathFlags(llvm_target->GetFastMathFlags());
cg->AddFunctionsOrdered(mod->functions.begin(), mod->functions.end());
- if (entry_func.length() != 0) {
- cg->AddMainFunction(entry_func);
- }
+ q
- module_owning_ptr_ = cg->Finish();
+ module_owning_ptr_ = cg->Finish();
module_ = module_owning_ptr_.get();
jit_engine_ = llvm_target->GetJITEngine();
llvm_target->SetTargetMetadata(module_);
diff --git a/src/target/source/codegen_c_host.cc b/src/target/source/codegen_c_host.cc
index 6cd12a9319..1c8a3dd2ea 100644
--- a/src/target/source/codegen_c_host.cc
+++ b/src/target/source/codegen_c_host.cc
@@ -72,20 +72,6 @@ void CodeGenCHost::AddFunction(const GlobalVar& gvar, const PrimFunc& func,
emit_fwd_func_decl_ = emit_fwd_func_decl;
CodeGenC::AddFunction(gvar, func);
- if (func->HasNonzeroAttr(tir::attr::kIsEntryFunc)) {
- ICHECK(global_symbol.has_value())
- << "CodeGenCHost: The entry func must have the global_symbol attribute, "
- << "but function " << gvar << " only has attributes " << func->attrs;
-
- function_names_.push_back(runtime::symbol::tvm_module_main);
- stream << "// CodegenC: NOTE: Auto-generated entry function\n";
- PrintFuncPrefix(stream);
- PrintType(func->ret_type, stream);
- stream << " " << tvm::runtime::symbol::tvm_module_main
- << "(void* self, void* args,int num_args, void* result) {\n";
- stream << " return " << global_symbol.value() << "(self, args, num_args, result);\n";
- stream << "}\n";
- }
}
void CodeGenCHost::GenerateForwardFunctionDeclarations(String global_symbol,
diff --git a/src/tir/ir/transform.cc b/src/tir/ir/transform.cc
index aafe6277e2..6ef7ffdca9 100644
--- a/src/tir/ir/transform.cc
+++ b/src/tir/ir/transform.cc
@@ -42,7 +42,6 @@ TVM_REGISTER_PASS_CONFIG_OPTION("tir.disable_cse_tir", Bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tir.enable_debug", Bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tir.enable_equiv_terms_in_cse_tir", Bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tir.disable_storage_rewrite", Bool);
-TVM_REGISTER_PASS_CONFIG_OPTION("tir.is_entry_func", Bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tir.add_lower_pass", Array<Array<ObjectRef>>);
TVM_REGISTER_PASS_CONFIG_OPTION("tir.debug_keep_trivial_loop", Bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tir.use_async_copy", Bool);
diff --git a/src/tir/transforms/primfunc_utils.cc b/src/tir/transforms/primfunc_utils.cc
index b1f3476eab..99cf901377 100644
--- a/src/tir/transforms/primfunc_utils.cc
+++ b/src/tir/transforms/primfunc_utils.cc
@@ -29,45 +29,6 @@ namespace tvm {
namespace tir {
namespace transform {
-transform::Pass AnnotateEntryFunc() {
- auto fpass = [](IRModule mod, transform::PassContext ctx) -> IRModule {
- // If only a single function exists, that function must be the entry
- if (mod->functions.size() == 1) {
- auto [gvar, base_func] = *mod->functions.begin();
- if (!base_func->HasNonzeroAttr(tir::attr::kIsEntryFunc)) {
- if (auto ptr = base_func.as<PrimFuncNode>()) {
- mod->Update(gvar, WithAttr(GetRef<PrimFunc>(ptr), tir::attr::kIsEntryFunc, true));
- }
- }
- return mod;
- }
-
- // If the module has multiple functions, but only one is exposed
- // externally, that function must be the entry.
- bool has_external_non_primfuncs = false;
- IRModule with_annotations;
- for (const auto& [gvar, base_func] : mod->functions) {
- bool is_external = base_func->GetAttr<String>(tvm::attr::kGlobalSymbol).has_value();
- if (is_external) {
- if (auto ptr = base_func.as<PrimFuncNode>()) {
- with_annotations->Add(gvar,
- WithAttr(GetRef<PrimFunc>(ptr), tir::attr::kIsEntryFunc, true));
- } else {
- has_external_non_primfuncs = true;
- }
- }
- }
- if (with_annotations->functions.size() == 1 && !has_external_non_primfuncs) {
- mod->Update(with_annotations);
- return mod;
- }
-
- // Default fallback, no annotations may be inferred.
- return mod;
- };
- return tvm::transform::CreateModulePass(fpass, 0, "tir.AnnotateEntryFunc", {});
-}
-
transform::Pass Filter(ffi::TypedFunction<bool(PrimFunc)> fcond) {
auto fpass = [fcond](tir::PrimFunc f, IRModule m, transform::PassContext ctx) {
if (fcond(f)) {
@@ -81,9 +42,7 @@ transform::Pass Filter(ffi::TypedFunction<bool(PrimFunc)> fcond) {
TVM_FFI_STATIC_INIT_BLOCK({
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef()
- .def("tir.transform.AnnotateEntryFunc", AnnotateEntryFunc)
- .def("tir.transform.Filter", Filter);
+ refl::GlobalDef().def("tir.transform.Filter", Filter);
});
} // namespace transform
diff --git a/tests/python/codegen/test_target_codegen_device.py b/tests/python/codegen/test_target_codegen_device.py
index 4dad03d700..0089e0bea6 100644
--- a/tests/python/codegen/test_target_codegen_device.py
+++ b/tests/python/codegen/test_target_codegen_device.py
@@ -95,7 +95,7 @@ def test_add_pipeline():
dev = tvm.device(device, 0)
target = tvm.target.Target(device, host)
mhost = tvm.tir.build(sch.mod, target=target)
- f = mhost.entry_func
+ f = mhost["main"]
# launch the kernel.
n = 1027
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
diff --git a/tests/python/contrib/test_hexagon/test_async_dma_pipeline.py b/tests/python/contrib/test_hexagon/test_async_dma_pipeline.py
index fe7a615531..808c016cf8 100644
--- a/tests/python/contrib/test_hexagon/test_async_dma_pipeline.py
+++ b/tests/python/contrib/test_hexagon/test_async_dma_pipeline.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
-""" Test different strategies for loading data into vtcm before running HVX workloads. """
+"""Test different strategies for loading data into vtcm before running HVX workloads."""
import numpy as np
import pytest
@@ -289,9 +289,13 @@ def evaluate(
if tvm.testing.utils.IS_IN_CI:
# Run with reduced number and repeat for CI
- timer = module.time_evaluator("__tvm_main__", hexagon_session.device, number=1, repeat=1)
+ timer = module.time_evaluator(
+ "__tvm_ffi_main__", hexagon_session.device, number=1, repeat=1
+ )
else:
- timer = module.time_evaluator("__tvm_main__", hexagon_session.device, number=10, repeat=10)
+ timer = module.time_evaluator(
+ "__tvm_ffi_main__", hexagon_session.device, number=10, repeat=10
+ )
time = timer(a_hexagon, b_hexagon, c_hexagon)
if expected_output is not None:
diff --git a/tests/python/contrib/test_hexagon/test_parallel_hvx.py b/tests/python/contrib/test_hexagon/test_parallel_hvx.py
index 8f77fa1c40..6822352568 100644
--- a/tests/python/contrib/test_hexagon/test_parallel_hvx.py
+++ b/tests/python/contrib/test_hexagon/test_parallel_hvx.py
@@ -160,7 +160,7 @@ def evaluate(hexagon_session, shape_dtypes, expected_output_producer, sch):
repeat = 1
timer = module.time_evaluator(
- "__tvm_main__", hexagon_session.device, number=number, repeat=repeat
+ "__tvm_ffi_main__", hexagon_session.device, number=number, repeat=repeat
)
runtime = timer(a_hexagon, b_hexagon, c_hexagon)
tvm.testing.assert_allclose(c_hexagon.numpy(), expected_output_producer(c_shape, a, b))
diff --git a/tests/python/contrib/test_hexagon/test_parallel_hvx_load_vtcm.py b/tests/python/contrib/test_hexagon/test_parallel_hvx_load_vtcm.py
index a584997dd5..63a65f3716 100644
--- a/tests/python/contrib/test_hexagon/test_parallel_hvx_load_vtcm.py
+++ b/tests/python/contrib/test_hexagon/test_parallel_hvx_load_vtcm.py
@@ -330,9 +330,7 @@ def setup_and_run(hexagon_session, sch, a, b, c, operations, mem_scope="global")
number = 1
repeat = 1
- timer = module.time_evaluator(
- "__tvm_main__", hexagon_session.device, number=number, repeat=repeat
- )
+ timer = module.time_evaluator("main", hexagon_session.device, number=number, repeat=repeat)
time = timer(a_hexagon, b_hexagon, c_hexagon)
gops = round(operations * 128 * 3 / time.mean / 1e9, 4)
return gops, c_hexagon.numpy()
@@ -364,9 +362,7 @@ def setup_and_run_preallocated(hexagon_session, sch, a, b, c, operations):
number = 1
repeat = 1
- timer = module.time_evaluator(
- "__tvm_main__", hexagon_session.device, number=number, repeat=repeat
- )
+ timer = module.time_evaluator("main", hexagon_session.device, number=number, repeat=repeat)
time = timer(a_hexagon, b_hexagon, c_hexagon, a_vtcm_hexagon, b_vtcm_hexagon, c_vtcm_hexagon)
gops = round(operations * 128 * 3 / time.mean / 1e9, 4)
return gops, c_hexagon.numpy()
diff --git a/tests/python/contrib/test_hexagon/test_parallel_scalar.py b/tests/python/contrib/test_hexagon/test_parallel_scalar.py
index bd9c78d5da..5c8043fdff 100644
--- a/tests/python/contrib/test_hexagon/test_parallel_scalar.py
+++ b/tests/python/contrib/test_hexagon/test_parallel_scalar.py
@@ -104,9 +104,7 @@ def evaluate(hexagon_session, operations, expected, sch):
number = 1
repeat = 1
- timer = module.time_evaluator(
- "__tvm_main__", hexagon_session.device, number=number, repeat=repeat
- )
+ timer = module.time_evaluator("main", hexagon_session.device, number=number, repeat=repeat)
runtime = timer(a_hexagon, b_hexagon, c_hexagon)
tvm.testing.assert_allclose(c_hexagon.numpy(), expected(a, b))
diff --git a/tests/python/contrib/test_hexagon/test_vtcm_bandwidth.py b/tests/python/contrib/test_hexagon/test_vtcm_bandwidth.py
index 931f99b2ec..265f2bf5fd 100644
--- a/tests/python/contrib/test_hexagon/test_vtcm_bandwidth.py
+++ b/tests/python/contrib/test_hexagon/test_vtcm_bandwidth.py
@@ -108,9 +108,13 @@ def evaluate(hexagon_session, sch, size):
if tvm.testing.utils.IS_IN_CI:
# Run with reduced number and repeat for CI
- timer = module.time_evaluator("__tvm_main__", hexagon_session.device, number=1, repeat=1)
+ timer = module.time_evaluator(
+ "__tvm_ffi_main__", hexagon_session.device, number=1, repeat=1
+ )
else:
- timer = module.time_evaluator("__tvm_main__", hexagon_session.device, number=10, repeat=10)
+ timer = module.time_evaluator(
+ "__tvm_ffi_main__", hexagon_session.device, number=10, repeat=10
+ )
runtime = timer(a_hexagon, a_vtcm_hexagon)
diff --git a/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py b/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py
index 81acf5ee86..0d2f445cb8 100644
--- a/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py
+++ b/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py
@@ -170,7 +170,7 @@ def set_global_func(head_dim, dtype):
with target:
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
f = tvm.tir.build(mod["main"], target=target)
- builts.append(f.entry_func)
+ builts.append(f["main"])
(
ftranspose_append,
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
index 1941edeaa7..305fd18f35 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
@@ -140,7 +140,7 @@ def set_global_func(head_dim, dtype):
with target:
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
f = tvm.tir.build(mod["main"], target=target)
- builts.append(f.entry_func)
+ builts.append(f["main"])
(
ftranspose_append,
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
index ffd3452292..e13ce1ca7b 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
@@ -156,7 +156,7 @@ def set_global_func():
with target:
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
f = tvm.tir.build(mod["main"], target=target)
- builts.append(f.entry_func)
+ builts.append(f["main"])
(
ftranspose_append,
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py
index 2f726064a7..53044a786c 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py
@@ -169,7 +169,7 @@ def set_global_func(dtype):
with target:
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
f = tvm.tir.build(mod["main"], target=target)
- builts.append(f.entry_func)
+ builts.append(f["main"])
(
ftranspose_append,
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py
index b2982abdb0..73a4d89dad 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py
@@ -134,7 +134,7 @@ def set_global_func(dtype):
with target:
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
f = tvm.tir.build(mod["main"], target=target)
- builts.append(f.entry_func)
+ builts.append(f["main"])
(
ftranspose_append,
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
index 8cd3a73740..44169828e2 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
@@ -142,7 +142,7 @@ def set_global_func(head_dim, dtype):
with target:
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
f = tvm.tir.build(mod["main"], target=target)
- builts.append(f.entry_func)
+ builts.append(f["main"])
(
ftranspose_append,
diff --git a/tests/python/relax/test_runtime_builtin_rnn_state.py b/tests/python/relax/test_runtime_builtin_rnn_state.py
index 095aba8b83..fe8c19257d 100644
--- a/tests/python/relax/test_runtime_builtin_rnn_state.py
+++ b/tests/python/relax/test_runtime_builtin_rnn_state.py
@@ -81,7 +81,7 @@ def set_global_func():
with target:
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod) # pylint: disable=not-callable
f = tvm.tir.build(mod["main"], target=target)
- return f.entry_func
+ return f["main"]
_f_tir_gets, _f_tir_sets = [], []
for state in states:
diff --git a/tests/python/tir-transform/test_tir_transform_helpers.py b/tests/python/tir-transform/test_tir_transform_helpers.py
index 0bbd0e7160..a7ac11d7a9 100644
--- a/tests/python/tir-transform/test_tir_transform_helpers.py
+++ b/tests/python/tir-transform/test_tir_transform_helpers.py
@@ -21,27 +21,6 @@ from tvm.script import tir as T, ir as I
import tvm.testing
-def test_annotate_entry_func_single_primfunc():
- @tvm.script.ir_module
- class MockModule:
- @T.prim_func(private=True)
- def func1(A: T.Buffer((16,), "float32")):
- for i in T.serial(16):
- if i == 5:
- if i == 5:
- A[i] = 0.0
-
- mod = MockModule
- assert mod
- assert not mod["func1"].attrs
- after = tvm.tir.transform.AnnotateEntryFunc()(mod)
- assert (
- after["func1"].attrs
- and "tir.is_entry_func" in after["func1"].attrs
- and after["func1"].attrs["tir.is_entry_func"]
- )
-
-
# Test module
@tvm.script.ir_module
class MockModule:
@@ -60,16 +39,6 @@ class MockModule:
A[i] = 0.0
-@pytest.mark.xfail
-def test_annotate_entry_func_multiple_primfunc():
- mod = MockModule
- assert mod
- assert not mod["func1"].attrs
- assert not mod["func2"].attrs
- # This should fail
- after = tvm.tir.transform.AnnotateEntryFunc()(mod)
-
-
def test_bind_target():
mod = MockModule
assert mod
diff --git a/tests/python/tvmscript/test_tvmscript_roundtrip.py b/tests/python/tvmscript/test_tvmscript_roundtrip.py
index 0e1b328844..73ca1dad3b 100644
--- a/tests/python/tvmscript/test_tvmscript_roundtrip.py
+++ b/tests/python/tvmscript/test_tvmscript_roundtrip.py
@@ -196,7 +196,6 @@ def opt_gemm_mod_host():
T.func_attr(
{
"tir.noalias": True,
- "tir.is_entry_func": True,
"calling_conv": 1,
}
)
@@ -2242,7 +2241,6 @@ def opt_conv_tensorcore_mod_host():
{
"tir.noalias": True,
"global_symbol": "default_function",
- "tir.is_entry_func": True,
"calling_conv": 1,
}
)