Fix TVMFFIEnvSetDLPackManagedTensorAllocator to correctly return the original allocator (#371)

Before this commit, TVMFFIEnvSetDLPackManagedTensorAllocator incorrectly
set the previous allocator in:
TVMFFIEnvSetDLPackManagedTensorAllocator(NewAllocator, 0,
&pre_allocator);

---------

Co-authored-by: nan <[email protected]>

Author: nan2088

src/ffi/extra/env_context.cc

@@ -65,13 +65,12 @@ class EnvContext {
   void SetDLPackManagedTensorAllocator(DLPackManagedTensorAllocator allocator,
                                        int write_to_global_context,
                                        DLPackManagedTensorAllocator* opt_out_original_allocator) {
-    dlpack_allocator_ = allocator;
+    if (opt_out_original_allocator != nullptr) {
+      *opt_out_original_allocator = GetDLPackManagedTensorAllocator();
+    }
     if (write_to_global_context != 0) {
       GlobalTensorAllocator() = allocator;
     }
-    if (opt_out_original_allocator != nullptr) {
-      *opt_out_original_allocator = dlpack_allocator_;
-    }
     dlpack_allocator_ = allocator;
   }
 

tests/cpp/extra/test_c_env_api.cc

@@ -59,6 +59,14 @@ int TestDLPackManagedTensorAllocatorError(DLTensor* prototype, DLManagedTensorVe
   return -1;
 }
 
+TEST(CEnvAPI, TVMFFIEnvSetDLPackManagedTensorAllocator) {
+  auto old_allocator = TVMFFIEnvGetDLPackManagedTensorAllocator();
+  DLPackManagedTensorAllocator pre_allocator;
+  TVMFFIEnvSetDLPackManagedTensorAllocator(TestDLPackManagedTensorAllocator, 0, &pre_allocator);
+  EXPECT_EQ(old_allocator, pre_allocator);
+  TVMFFIEnvSetDLPackManagedTensorAllocator(old_allocator, 0, nullptr);
+}
+
 TEST(CEnvAPI, TVMFFIEnvTensorAlloc) {
   auto old_allocator = TVMFFIEnvGetDLPackManagedTensorAllocator();
   TVMFFIEnvSetDLPackManagedTensorAllocator(TestDLPackManagedTensorAllocator, 0, nullptr);