Fix up the macro related reg

Author: tqchen

src/runtime/contrib/sort/sort.cc

@@ -80,7 +80,7 @@ struct float16 {
 // If input tensor has dimension (d0, d1, ..., d(k-1), dk, d(k+1), ..., d(n-1))
 // and sort axis is dk. sort_num should have dimension of
 // (d1, d2, ..., d(k-1), d(k+1), ..., dn).
-TVM_FFI_STATIC_INIT_BLOCK({
+void RegisterArgsortNMS() {
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef().def_packed(
       "tvm.contrib.sort.argsort_nms", [](ffi::PackedArgs args, ffi::Any* ret) {
@@ -157,7 +157,7 @@ TVM_FFI_STATIC_INIT_BLOCK({
           }
         }
       });
-});
+}
 
 template <typename DataType, typename OutType>
 void sort_impl(
@@ -222,7 +222,7 @@ void sort(DLTensor* input, DLTensor* output, int32_t axis, bool is_ascend) {
 // If input tensor has dimension (d0, d1, ..., d(k-1), dk, d(k+1), ..., d(n-1))
 // and sort axis is dk. sort_num should have dimension of
 // (d1, d2, ..., d(k-1), d(k+1), ..., dn).
-TVM_FFI_STATIC_INIT_BLOCK({
+void RegisterArgsort() {
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef().def_packed("tvm.contrib.sort.argsort", [](ffi::PackedArgs args, ffi::Any* ret) {
     auto input = args[0].cast<DLTensor*>();
@@ -311,7 +311,7 @@ TVM_FFI_STATIC_INIT_BLOCK({
       LOG(FATAL) << "Unsupported input dtype: " << data_dtype;
     }
   });
-});
+}
 
 // Sort implemented C library sort.
 // Return  sorted tensor.
@@ -320,7 +320,7 @@ TVM_FFI_STATIC_INIT_BLOCK({
 // If input tensor has dimension (d0, d1, ..., d(k-1), dk, d(k+1), ..., d(n-1))
 // and sort axis is dk. sort_num should have dimension of
 // (d1, d2, ..., d(k-1), d(k+1), ..., dn).
-TVM_FFI_STATIC_INIT_BLOCK({
+void RegisterSort() {
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef().def_packed("tvm.contrib.sort.sort", [](ffi::PackedArgs args, ffi::Any* ret) {
     auto input = args[0].cast<DLTensor*>();
@@ -357,7 +357,7 @@ TVM_FFI_STATIC_INIT_BLOCK({
       LOG(FATAL) << "Unsupported input dtype: " << data_dtype;
     }
   });
-});
+}
 
 template <typename DataType, typename IndicesType>
 void topk(DLTensor* input, DLTensor* out_values, DLTensor* out_indices, int k, int axis,
@@ -452,7 +452,7 @@ void topk(DLTensor* input, DLTensor* out_values, DLTensor* out_indices, int k, i
 // If input tensor has dimension (d0, d1, ..., d(k-1), dk, d(k+1), ..., d(n-1))
 // and sort axis is dk. sort_num should have dimension of
 // (d1, d2, ..., d(k-1), d(k+1), ..., dn).
-TVM_FFI_STATIC_INIT_BLOCK({
+void RegisterTopk() {
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef().def_packed("tvm.contrib.sort.topk", [](ffi::PackedArgs args, ffi::Any* ret) {
     auto input = args[0].cast<DLTensor*>();
@@ -574,6 +574,13 @@ TVM_FFI_STATIC_INIT_BLOCK({
       LOG(FATAL) << "Unsupported input dtype: " << data_dtype;
     }
   });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK({
+  RegisterArgsortNMS();
+  RegisterArgsort();
+  RegisterSort();
+  RegisterTopk();
 });
 
 }  // namespace contrib

src/target/llvm/codegen_llvm.cc

@@ -2351,7 +2351,7 @@ llvm::DIType* CodeGenLLVM::GetDebugType(const Type& ty_tir, llvm::Type* ty_llvm)
   return nullptr;
 }
 
-TVM_FFI_STATIC_INIT_BLOCK({
+static void CodegenLLVMRegisterReflection() {
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef()
       .def("tvm.codegen.llvm.GetDefaultTargetTriple",
@@ -2385,7 +2385,9 @@ TVM_FFI_STATIC_INIT_BLOCK({
         LOG(WARNING) << "Current version of LLVM does not support feature detection on your CPU";
         return {};
       });
-});
+}
+
+TVM_FFI_STATIC_INIT_BLOCK({ CodegenLLVMRegisterReflection(); });
 
 }  // namespace codegen
 }  // namespace tvm

src/target/llvm/llvm_module.cc

@@ -621,7 +621,7 @@ void* LLVMModuleNode::GetFunctionAddr(const std::string& name,
   return nullptr;
 }
 
-TVM_FFI_STATIC_INIT_BLOCK({
+static void LLVMReflectionRegister() {
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef()
       .def("target.build.llvm",
@@ -787,7 +787,9 @@ TVM_FFI_STATIC_INIT_BLOCK({
              n->SetJITEngine(llvm_target->GetJITEngine());
              return runtime::Module(n);
            });
-});
+}
+
+TVM_FFI_STATIC_INIT_BLOCK({ LLVMReflectionRegister(); });
 
 }  // namespace codegen
 }  // namespace tvm

tests/python/relax/test_vm_callback_function.py

@@ -100,6 +100,7 @@ def relax_func(
     )
     vm = tvm.relax.VirtualMachine(ex, dev)
 
+    # custom callback that raises an error in python
     def custom_callback():
         local_var = 42
         raise RuntimeError("Error thrown from callback")