(tvm-ffi) branch main updated: [Rust] Borrow BigInt, DLDataType and DLDevice in place with downcast_ref (#936)

[bohan.apache.org via commits]

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

spectrometerHBH 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 d37f62db [Rust] Borrow BigInt, DLDataType and DLDevice in place with 
downcast_ref (#936)
d37f62db is described below

commit d37f62db8940698d70dc6860ab08c2a79eba33f4
Author: Tianqi Chen <tqc...@users.noreply.github.com>
AuthorDate: Thu Nov 5 20:22:09 2026 -0500

    [Rust] Borrow BigInt, DLDataType and DLDevice in place with downcast_ref 
(#936)
    
    `downcast_ref` and `Array` indexing now borrow `DLDataType`, `DLDevice`
    and `BigInt` in place, as they do `i64` and `f64`. `DLDataType` and
    `DLDevice` are stored in the cell's union as themselves, and a `BigInt`
    is a whole cell, like `Any`.
    
    ```rust
    let dtype: &DLDataType = value.downcast_ref::<DLDataType>().unwrap();
    let n: &BigInt = value.downcast_ref::<BigInt>().unwrap();
    ```
    
    These two commits were pushed to #935 after it was reviewed and did not
    land with it.
---
 rust/tvm-ffi/src/type_traits.rs | 41 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 40 insertions(+), 1 deletion(-)

diff --git a/rust/tvm-ffi/src/type_traits.rs b/rust/tvm-ffi/src/type_traits.rs
index 56b9cd8c..8f8be8e2 100644
--- a/rust/tvm-ffi/src/type_traits.rs
+++ b/rust/tvm-ffi/src/type_traits.rs
@@ -244,6 +244,22 @@ unsafe impl RefFromAnyUnchecked for Any {
     }
 }
 
+// A `BigInt` is one cell, holding an inline `i64` or a `BigInt` object.
+unsafe impl RefFromAnyUnchecked for crate::big_int::BigInt {
+    #[inline(always)]
+    unsafe fn ref_from_any_unchecked(cell: &TVMFFIAny) -> &Self {
+        const {
+            assert!(
+                std::mem::size_of::<crate::big_int::BigInt>() == 
std::mem::size_of::<TVMFFIAny>()
+            );
+            assert!(
+                std::mem::align_of::<crate::big_int::BigInt>() == 
std::mem::align_of::<TVMFFIAny>()
+            );
+        };
+        &*(cell as *const TVMFFIAny as *const crate::big_int::BigInt)
+    }
+}
+
 unsafe impl RefFromAnyUnchecked for crate::String {
     #[inline(always)]
     unsafe fn ref_from_any_unchecked(cell: &TVMFFIAny) -> &Self {
@@ -273,6 +289,20 @@ unsafe impl RefFromAnyUnchecked for f64 {
     }
 }
 
+unsafe impl RefFromAnyUnchecked for tvm_ffi_sys::dlpack::DLDataType {
+    #[inline(always)]
+    unsafe fn ref_from_any_unchecked(cell: &TVMFFIAny) -> &Self {
+        &cell.data_union.v_dtype
+    }
+}
+
+unsafe impl RefFromAnyUnchecked for tvm_ffi_sys::dlpack::DLDevice {
+    #[inline(always)]
+    unsafe fn ref_from_any_unchecked(cell: &TVMFFIAny) -> &Self {
+        &cell.data_union.v_device
+    }
+}
+
 /// Types that [`crate::Any::downcast_ref`] and 
[`crate::AnyView::downcast_ref`]
 /// can borrow from an FFI cell in place: object types (`ArrayObj`), and the
 /// [`RefFromAnyUnchecked`] types (managed references such as `Array<T>`, 
`String`,
@@ -401,7 +431,16 @@ macro_rules! impl_downcast_ref_from_any {
     };
 }
 
-impl_downcast_ref_from_any!(Any, crate::String, crate::string::Bytes, i64, 
f64);
+impl_downcast_ref_from_any!(
+    Any,
+    crate::String,
+    crate::string::Bytes,
+    i64,
+    f64,
+    crate::big_int::BigInt,
+    tvm_ffi_sys::dlpack::DLDataType,
+    tvm_ffi_sys::dlpack::DLDevice,
+);
 
 mod container_element_ops {
     use super::{Any, AnyCompatible, AnyView, TVMFFIAny};