fix sort

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,10 +311,19 @@ TVM_FFI_STATIC_INIT_BLOCK({
       LOG(FATAL) << "Unsupported input dtype: " << data_dtype;
     }
   });
-});
+};
 
 
-void SortPacked(ffi::PackedArgs args, ffi::Any* ret) {
+// Sort implemented C library sort.
+// Return  sorted tensor.
+// By default, the last axis will be used to sort.
+// sort_num specify the number of elements to be sorted.
+// 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).
+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*>();
     auto output = args[1].cast<DLTensor*>();
     int32_t axis = args[2].cast<int32_t>();
@@ -348,18 +357,8 @@ void SortPacked(ffi::PackedArgs args, ffi::Any* ret) {
     } else {
       LOG(FATAL) << "Unsupported input dtype: " << data_dtype;
     }
+  });
 }
-// Sort implemented C library sort.
-// Return  sorted tensor.
-// By default, the last axis will be used to sort.
-// sort_num specify the number of elements to be sorted.
-// 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({
-  namespace refl = tvm::ffi::reflection;
-  refl::GlobalDef().def_packed("tvm.contrib.sort.sort", SortPacked);
-});
 
 template <typename DataType, typename IndicesType>
 void topk(DLTensor* input, DLTensor* out_values, DLTensor* out_indices, int k, int axis,
@@ -454,7 +453,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*>();
@@ -576,6 +575,13 @@ TVM_FFI_STATIC_INIT_BLOCK({
       LOG(FATAL) << "Unsupported input dtype: " << data_dtype;
     }
   });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK({
+  RegisterArgsortNMS();
+  RegisterArgsort();
+  RegisterSort();
+  RegisterTopk();
 });
 
 }  // namespace contrib

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