(tvm-ffi) branch main updated: [EXTRA] Improve the Error propagation in alloc (#224)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git


The following commit(s) were added to refs/heads/main by this push:
     new 227bdd0  [EXTRA] Improve the Error propagation in alloc (#224)
227bdd0 is described below

commit 227bdd0c5c70f186fac3b3c99427a032424ed58e
Author: Tianqi Chen <tqc...@users.noreply.github.com>
AuthorDate: Tue Nov 4 15:28:03 2025 -0500

    [EXTRA] Improve the Error propagation in alloc (#224)
    
    This PR updates the env_conetxt.cc to check the ret code first before
    checking the dl_tensor being not null, this allows us to effectively
    propagate error message from the caller if any to the consumer.
    Test case updated to guard this regression.
---
 python/tvm_ffi/error.py                                |  1 +
 python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py |  2 +-
 src/ffi/extra/env_context.cc                           |  4 +++-
 tests/cpp/extra/test_c_env_api.cc                      | 14 ++++++++++----
 4 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/python/tvm_ffi/error.py b/python/tvm_ffi/error.py
index 015fe80..03dfcca 100644
--- a/python/tvm_ffi/error.py
+++ b/python/tvm_ffi/error.py
@@ -214,3 +214,4 @@ register_error("AttributeError", AttributeError)
 register_error("KeyError", KeyError)
 register_error("IndexError", IndexError)
 register_error("AssertionError", AssertionError)
+register_error("MemoryError", MemoryError)
diff --git a/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py 
b/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py
index 77900ff..288b3ed 100644
--- a/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py
+++ b/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py
@@ -531,7 +531,7 @@ struct TorchDLPackExchangeAPI : public DLPackExchangeAPI {
       *out = at::toDLPackImpl<DLManagedTensorVersioned>(tensor);
       return 0;
     } catch (const std::exception& e) {
-      SetError(error_ctx, "TorchDLPackManagedTensorAllocator", e.what());
+      SetError(error_ctx, "MemoryError", e.what());
       return -1;
     }
   }
diff --git a/src/ffi/extra/env_context.cc b/src/ffi/extra/env_context.cc
index e8f5a63..cde76c4 100644
--- a/src/ffi/extra/env_context.cc
+++ b/src/ffi/extra/env_context.cc
@@ -139,8 +139,10 @@ int TVMFFIEnvTensorAlloc(DLTensor* prototype, 
TVMFFIObjectHandle* out) {
   }
   DLManagedTensorVersioned* dlpack_tensor = nullptr;
   int ret = (*dlpack_alloc)(prototype, &dlpack_tensor, nullptr, 
TVMFFIEnvTensorAllocSetError);
-  TVM_FFI_ICHECK(dlpack_tensor != nullptr);
   if (ret != 0) return ret;
+  // need to first check ret so we can properly propagate the error from 
caller side.
+  TVM_FFI_CHECK(dlpack_tensor != nullptr, RuntimeError)
+      << "TVMFFIEnvTensorAlloc: failed to allocate a tensor from the 
allocator";
   if (dlpack_tensor->dl_tensor.strides != nullptr || 
dlpack_tensor->dl_tensor.ndim == 0) {
     *out = tvm::ffi::details::ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr(
         
tvm::ffi::make_object<tvm::ffi::details::TensorObjFromDLPack<DLManagedTensorVersioned>>(
diff --git a/tests/cpp/extra/test_c_env_api.cc 
b/tests/cpp/extra/test_c_env_api.cc
index ff8a845..9ac51e3 100644
--- a/tests/cpp/extra/test_c_env_api.cc
+++ b/tests/cpp/extra/test_c_env_api.cc
@@ -48,7 +48,7 @@ int TestDLPackManagedTensorAllocatorError(DLTensor* 
prototype, DLManagedTensorVe
                                           void* error_ctx,
                                           void (*SetError)(void* error_ctx, 
const char* kind,
                                                            const char* 
message)) {
-  SetError(error_ctx, "RuntimeError", "TestDLPackManagedTensorAllocatorError");
+  SetError(error_ctx, "MemoryError", "TestDLPackManagedTensorAllocatorError");
   return -1;
 }
 
@@ -74,13 +74,19 @@ TEST(CEnvAPI, TVMFFIEnvTensorAlloc) {
 TEST(CEnvAPI, TVMFFIEnvTensorAllocError) {
   auto old_allocator = TVMFFIEnvGetDLPackManagedTensorAllocator();
   
TVMFFIEnvSetDLPackManagedTensorAllocator(TestDLPackManagedTensorAllocatorError, 
0, nullptr);
+
   EXPECT_THROW(
       {
-        Tensor::FromEnvAlloc(TVMFFIEnvTensorAlloc, {1, 2, 3}, 
DLDataType({kDLFloat, 32, 1}),
-                             DLDevice({kDLCPU, 0}));
+        try {
+          Tensor::FromEnvAlloc(TVMFFIEnvTensorAlloc, {1, 2, 3}, 
DLDataType({kDLFloat, 32, 1}),
+                               DLDevice({kDLCPU, 0}));
+        } catch (const tvm::ffi::Error& e) {
+          EXPECT_EQ(e.kind(), "MemoryError");
+          EXPECT_EQ(e.message(), "TestDLPackManagedTensorAllocatorError");
+          throw;
+        }
       },
       tvm::ffi::Error);
   TVMFFIEnvSetDLPackManagedTensorAllocator(old_allocator, 0, nullptr);
 }
-
 }  // namespace