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 9c834dad [PERF][Rust] Keep structural traversal fast paths in
registers (#903)
9c834dad is described below
commit 9c834dadb9d575be3818758f371d143720a1602a
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Nov 3 16:22:50 2026 -0500
[PERF][Rust] Keep structural traversal fast paths in registers (#903)
Keep structural traversal results as owning `TVMFFIAny` carriers along
native visitor and mutation paths, with `_raw` helpers and outlined
reflected, packed, and error handling. Construct slow-path aggregate
arguments inside those helpers and use private `C-unwind` return
boundaries so the AArch64 C convention can return the raw pair in
registers while preserving Rust cleanup. Keep lazy initialization
scratch inside outlined slow paths while ready loads stay inline.
Preserve public ABI, dispatcher guards, copy-on-write permission,
remapping, and error/interrupt ownership.
Always-inline lightweight internal traversal adapters and
ownership-transfer/drop wrappers while leaving generic one-shot public
entrypoints to compiler heuristics. Move visit context annotation into
`StructuralVisitEngine::with_visit_context_raw` and name the owning
result conversion `visit_result_from_any`.
Use `core::hint::cold_path()` directly at uncommon branches. This
requires Rust 1.95 or newer.
---
rust/tvm-ffi-macros/src/object_macros.rs | 26 ++-
rust/tvm-ffi/src/any.rs | 16 +-
rust/tvm-ffi/src/extra/structural_common.rs | 5 +-
rust/tvm-ffi/src/extra/structural_map.rs | 29 +--
rust/tvm-ffi/src/extra/structural_mutate.rs | 241 +++++++++++++++++++-----
rust/tvm-ffi/src/extra/structural_visit.rs | 282 ++++++++++++++++++++--------
rust/tvm-ffi/src/extra/structural_walk.rs | 81 +++++---
rust/tvm-ffi/src/object.rs | 5 +-
rust/tvm-ffi/src/reflection.rs | 1 +
9 files changed, 512 insertions(+), 174 deletions(-)
diff --git a/rust/tvm-ffi-macros/src/object_macros.rs
b/rust/tvm-ffi-macros/src/object_macros.rs
index 60120dd5..2bbf0714 100644
--- a/rust/tvm-ffi-macros/src/object_macros.rs
+++ b/rust/tvm-ffi-macros/src/object_macros.rs
@@ -165,7 +165,18 @@ pub(crate) fn expand_object(
tindex
}
);
- *TYPE_INDEX
+ // Keep Once's initialization descriptors in the cold
helper,
+ // rather than in every inlined runtime type-index caller.
+ #[cold]
+ #[inline(never)]
+ fn initialize_type_index() -> i32 {
+ *std::sync::LazyLock::force(&TYPE_INDEX)
+ }
+ if let Some(index) = std::sync::LazyLock::get(&TYPE_INDEX)
{
+ *index
+ } else {
+ initialize_type_index()
+ }
}
}
}
@@ -191,7 +202,18 @@ pub(crate) fn expand_object(
tindex
}
);
- *TYPE_INDEX
+ // Keep Once's initialization descriptors in the cold
helper,
+ // rather than in every inlined runtime type-index caller.
+ #[cold]
+ #[inline(never)]
+ fn initialize_type_index() -> i32 {
+ *std::sync::LazyLock::force(&TYPE_INDEX)
+ }
+ if let Some(index) = std::sync::LazyLock::get(&TYPE_INDEX)
{
+ *index
+ } else {
+ initialize_type_index()
+ }
}
}
}
diff --git a/rust/tvm-ffi/src/any.rs b/rust/tvm-ffi/src/any.rs
index feca0169..a1f2ea68 100644
--- a/rust/tvm-ffi/src/any.rs
+++ b/rust/tvm-ffi/src/any.rs
@@ -218,16 +218,17 @@ impl<T> TryFromTemp<T> {
// Any
//---------------------
impl Any {
+ #[inline(always)]
pub fn new() -> Self {
Self {
data: TVMFFIAny::new(),
}
}
- #[inline]
+ #[inline(always)]
pub fn type_index(&self) -> i32 {
self.data.type_index
}
- #[inline]
+ #[inline(always)]
pub(crate) fn as_raw_ffi_any(&self) -> &TVMFFIAny {
&self.data
}
@@ -295,18 +296,18 @@ impl Any {
}
}
- #[inline]
+ #[inline(always)]
pub unsafe fn as_data_ptr(&mut self) -> *mut TVMFFIAny {
&mut self.data
}
- #[inline]
+ #[inline(always)]
pub unsafe fn into_raw_ffi_any(this: Self) -> TVMFFIAny {
let this = std::mem::ManuallyDrop::new(this);
this.data
}
- #[inline]
+ #[inline(always)]
pub unsafe fn from_raw_ffi_any(data: TVMFFIAny) -> Self {
Self { data }
}
@@ -326,6 +327,7 @@ impl Any {
}
impl Default for Any {
+ #[inline(always)]
fn default() -> Self {
Self::new()
}
@@ -342,7 +344,7 @@ impl Clone for Any {
}
impl Drop for Any {
- #[inline]
+ #[inline(always)]
fn drop(&mut self) {
if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
unsafe { object::unsafe_::dec_ref(self.data.data_union.v_obj) }
@@ -402,7 +404,7 @@ fn any_view_to_owned_via_runtime(view: TVMFFIAny) -> Any {
}
impl<T: AnyCompatible> From<T> for Any {
- #[inline]
+ #[inline(always)]
fn from(value: T) -> Self {
unsafe {
let mut data = TVMFFIAny::new();
diff --git a/rust/tvm-ffi/src/extra/structural_common.rs
b/rust/tvm-ffi/src/extra/structural_common.rs
index 5979deb2..a7d36943 100644
--- a/rust/tvm-ffi/src/extra/structural_common.rs
+++ b/rust/tvm-ffi/src/extra/structural_common.rs
@@ -25,6 +25,7 @@ use crate::reflection::FieldGetter;
use crate::tvm_ffi_sys::{TVMFFIAny, TVMFFIByteArray, TVMFFITypeIndex,
TVMFFITypeKeyToIndex};
/// Add one structural traversal frame to an error's backtrace.
+#[cold]
pub(crate) fn with_structural_error_context(error: Error, operation: &str,
frame: &str) -> Error {
Error::with_appended_backtrace(error, &format!("[native structural
{operation}] {frame}\n"))
}
@@ -150,7 +151,7 @@ pub(crate) use impl_callback_chain_tuple_arities;
/// Raw string/byte views and `ObjectRValueRef` return `None` because they need
/// the runtime's normalization or move logic. A null object pointer also
/// returns `None` instead of constructing an invalid owning value.
-#[inline]
+#[inline(always)]
pub(crate) fn try_to_owned_without_normalization(raw: TVMFFIAny) ->
Option<Any> {
if is_plain_inline(raw.type_index) {
return Some(unsafe { Any::from_raw_ffi_any(raw) });
@@ -168,7 +169,7 @@ pub(crate) fn try_to_owned_without_normalization(raw:
TVMFFIAny) -> Option<Any>
pub(crate) use crate::any::is_plain_inline;
-#[inline]
+#[inline(always)]
pub(crate) fn same_shallow(lhs: TVMFFIAny, rhs: TVMFFIAny) -> bool {
lhs.type_index == rhs.type_index
&& lhs.small_str_len == rhs.small_str_len
diff --git a/rust/tvm-ffi/src/extra/structural_map.rs
b/rust/tvm-ffi/src/extra/structural_map.rs
index 4f965e29..c4778d18 100644
--- a/rust/tvm-ffi/src/extra/structural_map.rs
+++ b/rust/tvm-ffi/src/extra/structural_map.rs
@@ -93,7 +93,7 @@ pub enum ByMapDispatch {}
impl<'a, V: MapDispatch> IntoMapper<ByMapDispatch> for &'a mut V {
type Mapper = DispatchMapper<'a, V>;
- #[inline]
+ #[inline(always)]
fn into_mapper(self) -> Self::Mapper {
DispatchMapper(self)
}
@@ -180,7 +180,7 @@ where
T: crate::type_traits::AnyCompatible,
O: IntoMutateResult,
{
- #[inline]
+ #[inline(always)]
fn try_map(
&mut self,
value: AnyView<'_>,
@@ -201,7 +201,7 @@ where
T: crate::type_traits::AnyCompatible,
O: IntoMutateResult,
{
- #[inline]
+ #[inline(always)]
fn try_map(
&mut self,
value: AnyView<'_>,
@@ -222,7 +222,7 @@ where
N: ObjectCore,
O: IntoMutateResult,
{
- #[inline]
+ #[inline(always)]
fn try_map(
&mut self,
value: AnyView<'_>,
@@ -243,7 +243,7 @@ where
N: ObjectCore,
O: IntoMutateResult,
{
- #[inline]
+ #[inline(always)]
fn try_map(
&mut self,
value: AnyView<'_>,
@@ -263,7 +263,7 @@ where
F: for<'a> FnMut(AnyView<'a>) -> O,
O: IntoMutateResult,
{
- #[inline]
+ #[inline(always)]
fn try_map(
&mut self,
value: AnyView<'_>,
@@ -281,7 +281,7 @@ where
F: for<'a> FnMut(AnyView<'a>, DefRegionKind) -> O,
O: IntoMutateResult,
{
- #[inline]
+ #[inline(always)]
fn try_map(
&mut self,
value: AnyView<'_>,
@@ -298,7 +298,7 @@ pub struct
ByMapChainLink<Markers>(PhantomData<fn(Markers)>);
pub enum ByMapDispatchLink {}
impl<V: MapDispatch> MapChainLink<ByMapDispatchLink> for &mut V {
- #[inline]
+ #[inline(always)]
fn try_map(
&mut self,
value: AnyView<'_>,
@@ -316,7 +316,7 @@ pub struct MapChain<Link, Marker> {
}
impl<Link, Marker> MapChain<Link, Marker> {
- #[inline]
+ #[inline(always)]
fn new(link: Link) -> Self {
MapChain {
link,
@@ -352,7 +352,7 @@ macro_rules! impl_map_chain_link {
where
$($F: MapChainLink<$M>,)+
{
- #[inline]
+ #[inline(always)]
fn try_map(
&mut self,
value: AnyView<'_>,
@@ -373,7 +373,7 @@ macro_rules! impl_map_chain_link {
{
type Mapper = MapChain<($($F,)+), ByMapChainLink<($($M,)+)>>;
- #[inline]
+ #[inline(always)]
fn into_mapper(self) -> Self::Mapper {
MapChain::new(self)
}
@@ -394,7 +394,7 @@ macro_rules! impl_bare_map_link {
{
type Mapper = MapChain<F, $marker<T>>;
- #[inline]
+ #[inline(always)]
fn into_mapper(self) -> Self::Mapper {
MapChain::new(self)
}
@@ -417,7 +417,7 @@ where
{
type Mapper = MapChain<F, ByMapCatchAll>;
- #[inline]
+ #[inline(always)]
fn into_mapper(self) -> Self::Mapper {
MapChain::new(self)
}
@@ -430,7 +430,7 @@ where
{
type Mapper = MapChain<F, ByMapCatchAllKind>;
- #[inline]
+ #[inline(always)]
fn into_mapper(self) -> Self::Mapper {
MapChain::new(self)
}
@@ -479,6 +479,7 @@ impl<'v, 'c, 'r, 'a, D: MapDispatch, const PRE_ORDER: bool>
}
impl<D: MapDispatch, const PRE_ORDER: bool> NativeMapper<'_, D, PRE_ORDER> {
+ #[inline(always)]
fn map_current(
&mut self,
value: MutationView<'_>,
diff --git a/rust/tvm-ffi/src/extra/structural_mutate.rs
b/rust/tvm-ffi/src/extra/structural_mutate.rs
index 14565427..c4e6f7a4 100644
--- a/rust/tvm-ffi/src/extra/structural_mutate.rs
+++ b/rust/tvm-ffi/src/extra/structural_mutate.rs
@@ -45,8 +45,9 @@ use crate::tvm_ffi_sys::TVMFFIFieldFlagBitMask::{
kTVMFFIFieldFlagBitMaskSEqHashIgnore,
kTVMFFIFieldFlagBitSetterIsFunctionObj,
};
use crate::tvm_ffi_sys::{
- TVMFFIAny, TVMFFIAnyViewToOwnedAny, TVMFFIByteArray, TVMFFIFieldInfo,
TVMFFIFieldSetter,
- TVMFFIFunctionCall, TVMFFIGetTypeInfo, TVMFFIObject, TVMFFITypeIndex,
TVMFFITypeKeyToIndex,
+ TVMFFIAny, TVMFFIAnyDataUnion, TVMFFIAnyViewToOwnedAny, TVMFFIByteArray,
TVMFFIFieldInfo,
+ TVMFFIFieldSetter, TVMFFIFunctionCall, TVMFFIGetTypeInfo, TVMFFIObject,
TVMFFITypeIndex,
+ TVMFFITypeKeyToIndex,
};
use crate::tvm_ffi_sys::{TVMFFIObjectHandle, TVMFFISEqHashKind};
@@ -157,6 +158,7 @@ struct MemoEntry {
}
impl Default for StructuralMutatorObj {
+ #[inline(always)]
fn default() -> Self {
Self {
base: Object::new(),
@@ -222,11 +224,13 @@ impl StructuralMutateEngine {
) -> *mut c_void {
let context = std::ptr::from_mut(dispatch).cast();
if context != (*Self::from_handle(runtime)).dispatcher {
+ core::hint::cold_path();
wrong_mutation_dispatcher();
}
context
}
+ #[inline(always)]
fn vtable<D, const CALLBACK_CONTEXT: bool>() -> &'static
StructuralMutatorVTable
where
for<'s, 'v, 'c, 'r> &'s mut D: Dispatcher<
@@ -273,13 +277,11 @@ impl StructuralMutateEngine {
) {
// Map adapters attribute errors to their possibly replaced value;
// direct mutation attributes callback failures to this input.
- Some(result) => {
- if CALLBACK_CONTEXT {
- result.map_err(|error| with_value_context(error, raw))
- } else {
- result
- }
- }
+ Some(result) => result_into_raw(if CALLBACK_CONTEXT {
+ result.map_err(|error| with_value_context(error, raw))
+ } else {
+ result
+ }),
None => Self::default_mutate_raw(
runtime,
std::ptr::from_mut(dispatch).cast(),
@@ -289,7 +291,7 @@ impl StructuralMutateEngine {
};
// All callback reborrows have ended before restoring the caller's
provenance.
(*Self::from_handle(runtime)).dispatcher = context;
- result_into_raw(result)
+ result
}
#[inline(always)]
@@ -298,13 +300,53 @@ impl StructuralMutateEngine {
context: *mut c_void,
raw: TVMFFIAny,
permit: Permit,
- ) -> Result<Any> {
+ ) -> TVMFFIAny {
(*Self::from_handle(runtime)).dispatcher = context;
let kind = std::mem::transmute::<i32,
DefRegionKind>((*runtime).def_region_mode);
- Self::mutate_children(runtime, raw, kind, permit)
- .map_err(|error| with_value_context(error, raw))
+ Self::mutate_children_raw(runtime, raw, kind, permit)
}
+ #[inline(always)]
+ unsafe fn with_native_mutation_context_raw(result: TVMFFIAny, raw:
TVMFFIAny) -> TVMFFIAny {
+ if result.type_index != TVMFFITypeIndex::kTVMFFIError as i32 {
+ result
+ } else {
+ core::hint::cold_path();
+ Self::with_mutation_context_raw(
+ result.type_index,
+ result.small_str_len,
+ result.data_union,
+ raw.type_index,
+ raw.small_str_len,
+ raw.data_union,
+ )
+ }
+ }
+
+ #[cold]
+ #[inline(never)]
+ unsafe extern "C-unwind" fn with_mutation_context_raw(
+ result_type_index: i32,
+ result_small_str_len: u32,
+ result_data: TVMFFIAnyDataUnion,
+ type_index: i32,
+ small_str_len: u32,
+ data_union: TVMFFIAnyDataUnion,
+ ) -> TVMFFIAny {
+ let result = TVMFFIAny {
+ type_index: result_type_index,
+ small_str_len: result_small_str_len,
+ data_union: result_data,
+ };
+ let raw = TVMFFIAny {
+ type_index,
+ small_str_len,
+ data_union,
+ };
+ result_into_raw(result_from_raw(result).map_err(|error|
with_value_context(error, raw)))
+ }
+
+ #[inline(always)]
unsafe fn runtime_var_remap_get(
runtime: StructuralMutatorHandle,
raw: TVMFFIAny,
@@ -315,6 +357,7 @@ impl StructuralMutateEngine {
.get(borrowed_view(&raw))
}
+ #[inline(always)]
unsafe fn runtime_var_remap_set(
runtime: StructuralMutatorHandle,
raw: TVMFFIAny,
@@ -488,6 +531,7 @@ impl StructuralMutateEngine {
*field_changed = true;
Ok(())
}
+ #[inline(always)]
fn call_mutator(
mutator: StructuralMutatorHandle,
raw: TVMFFIAny,
@@ -507,11 +551,11 @@ impl StructuralMutateEngine {
}
}
- #[inline]
+ #[inline(always)]
fn structural_mutate_hook(raw: TVMFFIAny, permit: Permit) ->
Option<TVMFFIAny> {
let use_inplace = permit == Permit::MaybeInPlace &&
object_is_unique(raw);
if use_inplace {
- if let Some(attr) =
STRUCTURAL_MAYBE_INPLACE_MUTATE_COLUMN.get_raw(raw.type_index) {
+ if let Some(attr) =
structural_mutate_column::<true>().get_raw(raw.type_index) {
if attr.type_index == TVMFFITypeIndex::kTVMFFIOpaquePtr as i32
|| attr.type_index == TVMFFITypeIndex::kTVMFFIFunction as
i32
{
@@ -519,50 +563,75 @@ impl StructuralMutateEngine {
}
}
}
- let attr = STRUCTURAL_MUTATE_COLUMN.get_raw(raw.type_index)?;
+ let attr =
structural_mutate_column::<false>().get_raw(raw.type_index)?;
(attr.type_index != TVMFFITypeIndex::kTVMFFINone as
i32).then_some(attr)
}
#[inline(always)]
- fn call_structural_mutate_hook(
+ fn call_structural_mutate_hook_raw(
mutator: StructuralMutatorHandle,
raw: TVMFFIAny,
attr: TVMFFIAny,
- ) -> Result<Any> {
+ ) -> TVMFFIAny {
unsafe {
if attr.type_index == TVMFFITypeIndex::kTVMFFIOpaquePtr as i32 {
- let pointer = attr.data_union.v_ptr;
- // The registration protocol defines this opaque pointer as
FStructuralMutate.
- let hook: FStructuralMutate = std::mem::transmute(pointer);
- result_from_raw(hook(mutator, AnyView::from_raw_ffi_any(raw)))
+ let hook: FStructuralMutate =
std::mem::transmute(attr.data_union.v_ptr);
+ Self::with_native_mutation_context_raw(
+ hook(mutator, AnyView::from_raw_ffi_any(raw)),
+ raw,
+ )
} else {
- Self::call_structural_mutate_hook_slow(mutator, raw, attr)
+ Self::call_structural_mutate_hook_slow_raw(
+ mutator,
+ raw.type_index,
+ raw.small_str_len,
+ raw.data_union,
+ attr.type_index,
+ attr.small_str_len,
+ attr.data_union,
+ )
}
}
}
- // Keep packed calls and malformed registrations out of the native-hook
hot path.
+ // Construct packed-call aggregates only in the outlined slow path.
#[cold]
- fn call_structural_mutate_hook_slow(
+ #[inline(never)]
+ unsafe extern "C-unwind" fn call_structural_mutate_hook_slow_raw(
mutator: StructuralMutatorHandle,
- raw: TVMFFIAny,
- attr: TVMFFIAny,
- ) -> Result<Any> {
- unsafe {
- if attr.type_index == TVMFFITypeIndex::kTVMFFIFunction as i32 {
- let function =
Function::try_from(AnyView::from_raw_ffi_any(attr))?;
+ type_index: i32,
+ small_str_len: u32,
+ data_union: TVMFFIAnyDataUnion,
+ attr_type_index: i32,
+ attr_small_str_len: u32,
+ attr_data_union: TVMFFIAnyDataUnion,
+ ) -> TVMFFIAny {
+ let raw = TVMFFIAny {
+ type_index,
+ small_str_len,
+ data_union,
+ };
+ let attr = TVMFFIAny {
+ type_index: attr_type_index,
+ small_str_len: attr_small_str_len,
+ data_union: attr_data_union,
+ };
+ let result = if attr.type_index == TVMFFITypeIndex::kTVMFFIFunction as
i32 {
+
Function::try_from(AnyView::from_raw_ffi_any(attr)).and_then(|function| {
function.call_packed(&[
Self::borrowed_mutator_view(mutator),
AnyView::from_raw_ffi_any(raw),
])
- } else {
- Err(Error::new(
- TYPE_ERROR,
- "__s_mutate__ must be an opaque function pointer or
ffi.Function",
- "",
- ))
- }
- }
+ })
+ } else {
+ Err(Error::new(
+ TYPE_ERROR,
+ "__s_mutate__ must be an opaque function pointer or
ffi.Function",
+ "",
+ ))
+ };
+ // Preserve successful Error-valued returns until the public
conversion.
+ result_into_raw(result.map_err(|error| with_value_context(error, raw)))
}
/// Borrow a live runtime mutator as an object-valued ABI argument.
@@ -571,6 +640,7 @@ impl StructuralMutateEngine {
///
/// `mutator` must point to a live object and outlive the returned view.
The
/// view does not increment the object's reference count.
+ #[inline(always)]
unsafe fn borrowed_mutator_view<'a>(mutator: StructuralMutatorHandle) ->
AnyView<'a> {
let object = mutator.cast::<TVMFFIObject>();
let mut raw = TVMFFIAny::new();
@@ -598,6 +668,7 @@ impl StructuralMutateEngine {
previous: i32,
}
impl Drop for Restore {
+ #[inline(always)]
fn drop(&mut self) {
unsafe { (*self.mutator).def_region_mode = self.previous };
}
@@ -608,18 +679,61 @@ impl StructuralMutateEngine {
}
}
- fn mutate_children(
+ // Operate in the raw path to optimize stack frame cost. Each returned raw
+ // carrier transfers one owning result; fallible/destructor-bearing code
+ // imports it into Any before performing work that may unwind.
+ #[inline(always)]
+ fn mutate_children_raw(
runtime: StructuralMutatorHandle,
raw: TVMFFIAny,
def_region_kind: DefRegionKind,
permit: Permit,
- ) -> Result<Any> {
- // Match C++ DefaultMutateExpected: a registered type hook owns any
- // identity-remap policy for that type. Automatic remapping applies
only
- // to the reflected fallback below.
- if let Some(mutated) = Self::structural_mutate_hook(raw, permit) {
- return Self::call_structural_mutate_hook(runtime, raw, mutated);
+ ) -> TVMFFIAny {
+ // Registered hooks own their remapping policy; reflection handles its
own.
+ if let Some(attr) = Self::structural_mutate_hook(raw, permit) {
+ return Self::call_structural_mutate_hook_raw(runtime, raw, attr);
+ }
+ unsafe {
+ Self::mutate_without_hook_raw(
+ runtime,
+ raw.type_index,
+ raw.small_str_len,
+ raw.data_union,
+ def_region_kind as i32,
+ )
}
+ }
+
+ // The AArch64 C ABI keeps raw results in registers; C-unwind preserves
+ // Rust unwinding. Reflection is a normal fallback, not universally cold.
+ #[inline(never)]
+ unsafe extern "C-unwind" fn mutate_without_hook_raw(
+ runtime: StructuralMutatorHandle,
+ type_index: i32,
+ small_str_len: u32,
+ data_union: TVMFFIAnyDataUnion,
+ def_region_kind: i32,
+ ) -> TVMFFIAny {
+ let raw = TVMFFIAny {
+ type_index,
+ small_str_len,
+ data_union,
+ };
+ result_into_raw(
+ Self::mutate_without_hook(
+ runtime,
+ raw,
+ std::mem::transmute::<i32, DefRegionKind>(def_region_kind),
+ )
+ .map_err(|error| with_value_context(error, raw)),
+ )
+ }
+
+ fn mutate_without_hook(
+ runtime: StructuralMutatorHandle,
+ raw: TVMFFIAny,
+ def_region_kind: DefRegionKind,
+ ) -> Result<Any> {
if raw.type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 {
return owned_from_raw(raw);
}
@@ -881,6 +995,30 @@ static STRUCTURAL_MAYBE_INPLACE_MUTATE_COLUMN:
LazyLock<TypeAttrColumn> = LazyLo
TypeAttrColumn::new(STRUCTURAL_MAYBE_INPLACE_MUTATE_ATTR)
.expect("missing __s_maybe_inplace_mutate__ column")
});
+#[inline(always)]
+fn structural_mutate_column<const INPLACE: bool>() -> TypeAttrColumn {
+ let column = if INPLACE {
+ &STRUCTURAL_MAYBE_INPLACE_MUTATE_COLUMN
+ } else {
+ &STRUCTURAL_MUTATE_COLUMN
+ };
+ if let Some(column) = LazyLock::get(column) {
+ *column
+ } else {
+ initialize_structural_mutate_column::<INPLACE>()
+ }
+}
+
+#[cold]
+#[inline(never)]
+fn initialize_structural_mutate_column<const INPLACE: bool>() ->
TypeAttrColumn {
+ if INPLACE {
+ *LazyLock::force(&STRUCTURAL_MAYBE_INPLACE_MUTATE_COLUMN)
+ } else {
+ *LazyLock::force(&STRUCTURAL_MUTATE_COLUMN)
+ }
+}
+
static SHALLOW_COPY_COLUMN: LazyLock<TypeAttrColumn> = LazyLock::new(|| {
TypeAttrColumn::new(SHALLOW_COPY_ATTR).expect("missing
__ffi_shallow_copy__ column")
});
@@ -917,14 +1055,14 @@ pub trait IntoMutateResult:
callback_result_sealed::Sealed {
}
impl<T: Into<Any>> IntoMutateResult for T {
- #[inline]
+ #[inline(always)]
fn into_mutate_result(self) -> Result<Any> {
Ok(self.into())
}
}
impl<T: Into<Any>> IntoMutateResult for Result<T> {
- #[inline]
+ #[inline(always)]
fn into_mutate_result(self) -> Result<Any> {
self.map(Into::into)
}
@@ -944,6 +1082,7 @@ pub enum InplaceMode {
}
impl InplaceMode {
+ #[inline(always)]
pub(super) fn permit(self) -> Permit {
match self {
Self::Disallow => Permit::Copy,
@@ -1254,12 +1393,12 @@ impl<D> StructuralMutatorView<'_, D> {
let value = value.into();
unsafe {
let context =
StructuralMutateEngine::checked_dispatcher(self.runtime, dispatch);
- StructuralMutateEngine::default_mutate_raw(
+ result_from_raw(StructuralMutateEngine::default_mutate_raw(
self.runtime,
context,
value.value,
value.permit(),
- )
+ ))
}
.and_then(UnchangedOr::from_carrier)
}
@@ -1292,6 +1431,7 @@ where
for<'s, 'v, 'c, 'r> &'s mut D:
Dispatcher<(MutationView<'v>, &'c mut StructuralMutatorView<'r, D>),
Option<Result<Any>>>,
{
+ #[inline(always)]
fn mutate_root(mut self, root: Any) -> Result<Any> {
StructuralMutateEngine::run_mutator::<D, true>(root, &mut self)
}
@@ -1302,6 +1442,7 @@ where
for<'s, 'v, 'c, 'r> &'s mut D:
Dispatcher<(MutationView<'v>, &'c mut StructuralMutatorView<'r, D>),
Option<Result<Any>>>,
{
+ #[inline(always)]
fn mutate_root(self, root: Any) -> Result<Any> {
StructuralMutateEngine::run_mutator::<D, true>(root, self)
}
diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs
b/rust/tvm-ffi/src/extra/structural_visit.rs
index 0cf2c204..e36d6428 100644
--- a/rust/tvm-ffi/src/extra/structural_visit.rs
+++ b/rust/tvm-ffi/src/extra/structural_visit.rs
@@ -67,8 +67,8 @@ use crate::tvm_ffi_sys::TVMFFIFieldFlagBitMask::{
kTVMFFIFieldFlagBitMaskSEqHashIgnore,
};
use crate::tvm_ffi_sys::{
- TVMFFIAny, TVMFFIByteArray, TVMFFIDefRegionKind, TVMFFIFieldInfo,
TVMFFIGetTypeInfo,
- TVMFFIObject, TVMFFISEqHashKind, TVMFFITypeIndex, TVMFFITypeKeyToIndex,
+ TVMFFIAny, TVMFFIAnyDataUnion, TVMFFIByteArray, TVMFFIDefRegionKind,
TVMFFIFieldInfo,
+ TVMFFIGetTypeInfo, TVMFFIObject, TVMFFISEqHashKind, TVMFFITypeIndex,
TVMFFITypeKeyToIndex,
};
use super::structural_common::{
@@ -155,6 +155,7 @@ pub(super) struct StructuralVisitEngine {
}
impl Default for StructuralVisitorObj {
+ #[inline(always)]
fn default() -> Self {
Self {
base: Object::new(),
@@ -206,6 +207,7 @@ crate::def_object_ref! {
impl VisitInterrupt {
/// Interrupt carrying an FFI-compatible payload.
+ #[inline(always)]
pub fn with<T: Into<Any>>(payload: T) -> Self {
VisitInterruptObj {
base: Object::new(),
@@ -219,7 +221,55 @@ impl VisitInterrupt {
// StructuralVisitEngine implementation
//
-----------------------------------------------------------------------------
+// Operate in the raw path to optimize stack frame cost. Each `_raw` result
+// carries one owned ABI value: forward it exactly once or import it into Any
+// before any operation that can unwind or discard it.
impl StructuralVisitEngine {
+ #[inline(always)]
+ fn with_visit_context_raw(result: TVMFFIAny, value: TVMFFIAny) ->
TVMFFIAny {
+ if result.type_index != TVMFFITypeIndex::kTVMFFIError as i32
+ || value.type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as
i32
+ {
+ result
+ } else {
+ core::hint::cold_path();
+ // SAFETY: the Error tag establishes the owned object's
representation.
+ unsafe {
+ let error = take_error(Any::from_raw_ffi_any(result));
+ Any::into_raw_ffi_any(
+ Self::annotate_visit_error(
+ error,
+ value.type_index,
+ value.small_str_len,
+ value.data_union,
+ )
+ .into(),
+ )
+ }
+ }
+ }
+
+ #[cold]
+ #[inline(never)]
+ fn annotate_visit_error(
+ error: Error,
+ type_index: i32,
+ small_str_len: u32,
+ data_union: TVMFFIAnyDataUnion,
+ ) -> Error {
+ let value = TVMFFIAny {
+ type_index,
+ small_str_len,
+ data_union,
+ };
+ let error = with_visit_error_context(error, value);
+ with_structural_error_context(
+ error,
+ "visit",
+ &format!("object `{}`", type_key_of(value.type_index)),
+ )
+ }
+
/// Recover the Rust engine whose first field is this ABI base.
///
/// # Safety
@@ -236,29 +286,48 @@ impl StructuralVisitEngine {
/// Descent itself never accesses D: every child reenters the same visitor
/// callback, which restores its incoming pointer before the next sibling.
#[inline(always)]
- unsafe fn visit_children(runtime: StructuralVisitorHandle, value:
TVMFFIAny) -> Any {
- if let Some(attr) = STRUCTURAL_VISIT_COLUMN.get_raw(value.type_index) {
+ unsafe fn visit_children_raw(runtime: StructuralVisitorHandle, value:
TVMFFIAny) -> TVMFFIAny {
+ if let Some(attr) =
structural_visit_column().get_raw(value.type_index) {
if attr.type_index != TVMFFITypeIndex::kTVMFFINone as i32 {
- return Self::call_structural_visit_hook(runtime, value, attr);
+ return Self::call_structural_visit_hook_raw(runtime, value,
attr);
}
}
if value.type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32
{
- Any::new()
+ TVMFFIAny::new()
} else {
- Self::visit_reflected_fields(runtime, value)
+ Self::visit_reflected_fields_raw(
+ runtime,
+ value.type_index,
+ value.small_str_len,
+ value.data_union,
+ )
}
}
+ // The AArch64 C ABI returns TVMFFIAny in registers; C-unwind preserves
+ // Rust cleanup if reflected traversal panics inside this private helper.
#[inline(never)]
- unsafe fn visit_reflected_fields(runtime: StructuralVisitorHandle, value:
TVMFFIAny) -> Any {
+ unsafe extern "C-unwind" fn visit_reflected_fields_raw(
+ runtime: StructuralVisitorHandle,
+ type_index: i32,
+ small_str_len: u32,
+ data_union: TVMFFIAnyDataUnion,
+ ) -> TVMFFIAny {
+ let value = TVMFFIAny {
+ type_index,
+ small_str_len,
+ data_union,
+ };
let type_info = unsafe { TVMFFIGetTypeInfo(value.type_index) };
if type_info.is_null() {
- return runtime_error(&format!(
- "native visitor: unregistered type index {}",
- value.type_index
- ))
- .into();
+ return Any::into_raw_ffi_any(
+ runtime_error(&format!(
+ "native visitor: unregistered type index {}",
+ value.type_index
+ ))
+ .into(),
+ );
}
let seq_hash_kind = unsafe {
let metadata = (*type_info).metadata;
@@ -281,7 +350,7 @@ impl StructuralVisitEngine {
}
})
};
- halted.unwrap_or_else(Any::new)
+ Any::into_raw_ffi_any(halted.unwrap_or_else(Any::new))
}
unsafe fn visit_reflected_field(
@@ -314,17 +383,21 @@ impl StructuralVisitEngine {
let child_region = field_def_region(field, inherited_region);
with_error_context(
Self::with_visitor_def_region(runtime, child_region, || {
- Self::call_visitor(runtime, raw_of(borrowed))
+ Any::from_raw_ffi_any(Self::call_visitor_raw(runtime,
raw_of(borrowed)))
}),
&format!("field `{}`", field.name.as_str()),
)
}
#[inline(always)]
- unsafe fn visit<D>(runtime: StructuralVisitorHandle, dispatch: &mut D,
raw: TVMFFIAny) -> Any {
+ unsafe fn visit_raw<D>(
+ runtime: StructuralVisitorHandle,
+ dispatch: &mut D,
+ raw: TVMFFIAny,
+ ) -> TVMFFIAny {
let context = Self::checked_dispatcher(runtime, dispatch);
(*Self::from_handle(runtime)).dispatcher = context;
- Self::call_visitor(runtime, raw)
+ Self::call_visitor_raw(runtime, raw)
}
// The caller has already established matching runtime/D types and scope.
@@ -336,11 +409,13 @@ impl StructuralVisitEngine {
) -> *mut c_void {
let current = dispatcher_context(dispatch);
if current != (*Self::from_handle(runtime)).dispatcher {
+ core::hint::cold_path();
wrong_visitor_dispatcher();
}
current
}
+ #[inline(always)]
fn vtable<D>() -> &'static StructuralVisitorVTable
where
for<'s, 'v, 'c, 'r> &'s mut D: crate::Dispatcher<
@@ -349,11 +424,11 @@ impl StructuralVisitEngine {
>,
{
&StructuralVisitorVTable {
- visit: Self::rust_vtable_visit::<D>,
+ visit: Self::rust_vtable_visit_raw::<D>,
}
}
- unsafe extern "C" fn rust_vtable_visit<D>(
+ unsafe extern "C" fn rust_vtable_visit_raw<D>(
runtime: StructuralVisitorHandle,
value: AnyView<'static>,
) -> TVMFFIAny
@@ -366,7 +441,8 @@ impl StructuralVisitEngine {
let context = (*Self::from_handle(runtime)).dispatcher;
let raw = *value.as_raw_ffi_any();
let result = if raw.type_index == TVMFFITypeIndex::kTVMFFINone as i32 {
- Any::new()
+ core::hint::cold_path();
+ TVMFFIAny::new()
} else {
let dispatch = &mut *context.cast::<D>();
let mut visitor = Self::visitor_context::<D>(runtime);
@@ -374,14 +450,14 @@ impl StructuralVisitEngine {
&mut *dispatch,
(borrowed_view(&raw), &mut visitor),
) {
- Some(result) => visit_result(result, raw),
+ Some(result) => visit_result_raw(result, raw),
None => Self::default_visit_raw(runtime,
dispatcher_context(dispatch), raw),
}
};
// The enclosing callback may resume using its ancestor mutable borrow.
// Native hooks cannot restore Rust-only borrow provenance for us.
(*Self::from_handle(runtime)).dispatcher = context;
- Any::into_raw_ffi_any(result)
+ result
}
/// # Safety
@@ -402,14 +478,14 @@ impl StructuralVisitEngine {
runtime: StructuralVisitorHandle,
context: *mut c_void,
raw: TVMFFIAny,
- ) -> Any {
+ ) -> TVMFFIAny {
// Context is the supplied exclusive reborrow. Publish it before
descent;
// the handle-only scanner never creates or uses a competing D
reference.
(*Self::from_handle(runtime)).dispatcher = context;
- with_value_context(Self::visit_children(runtime, raw), raw)
+ Self::with_visit_context_raw(Self::visit_children_raw(runtime, raw),
raw)
}
- pub(super) fn run_visitor<D>(root: TVMFFIAny, dispatch: &mut D) -> Any
+ pub(super) fn run_visitor_raw<D>(root: TVMFFIAny, dispatch: &mut D) ->
TVMFFIAny
where
for<'s, 'v, 'c, 'r> &'s mut D: crate::Dispatcher<
(AnyView<'v>, &'c mut StructuralVisitorView<'r, D>),
@@ -435,49 +511,78 @@ impl StructuralVisitEngine {
(*handle).vtable = Self::vtable::<D>();
(*handle).def_region_mode = DefRegionKind::None as i32;
}
- let result = Self::call_visitor(handle, root);
+ // Keep the result owned while releasing the runtime, including on
unwind.
+ let result = unsafe {
Any::from_raw_ffi_any(Self::call_visitor_raw(handle, root)) };
// Trusted hooks have released every temporary strong/weak runtime
reference.
drop(engine);
- result
+ unsafe { Any::into_raw_ffi_any(result) }
}
#[inline(always)]
- fn call_visitor(visitor: StructuralVisitorHandle, raw: TVMFFIAny) -> Any {
+ fn call_visitor_raw(visitor: StructuralVisitorHandle, raw: TVMFFIAny) ->
TVMFFIAny {
if raw.type_index == TVMFFITypeIndex::kTVMFFINone as i32 {
- return Any::new();
+ core::hint::cold_path();
+ return TVMFFIAny::new();
}
unsafe {
let callback = (*(*visitor).vtable).visit;
- Any::from_raw_ffi_any(callback(visitor,
AnyView::from_raw_ffi_any(raw)))
+ callback(visitor, AnyView::from_raw_ffi_any(raw))
}
}
#[inline(always)]
- fn call_structural_visit_hook(
+ fn call_structural_visit_hook_raw(
visitor: StructuralVisitorHandle,
raw: TVMFFIAny,
attr: TVMFFIAny,
- ) -> Any {
+ ) -> TVMFFIAny {
unsafe {
if attr.type_index == TVMFFITypeIndex::kTVMFFIOpaquePtr as i32 {
let pointer = attr.data_union.v_ptr;
// Native pointer registration is an unsafe contract: a
non-null
// callback with this ABI, borrowed runtime lifetime and
owning result.
let hook: FStructuralVisit = std::mem::transmute(pointer);
- Any::from_raw_ffi_any(hook(visitor,
AnyView::from_raw_ffi_any(raw)))
+ hook(visitor, AnyView::from_raw_ffi_any(raw))
} else {
- Self::call_structural_visit_hook_slow(visitor, raw, attr)
+ core::hint::cold_path();
+ Self::call_structural_visit_hook_slow_raw(
+ visitor,
+ raw.type_index,
+ raw.small_str_len,
+ raw.data_union,
+ attr.type_index,
+ attr.small_str_len,
+ attr.data_union,
+ )
}
}
}
// Keep packed calls and malformed registrations out of the native-hook
hot path.
+ // Match the native raw-result ABI (registers on AArch64), without adding
+ // an abort boundary around Rust packed-call cleanup.
#[cold]
- fn call_structural_visit_hook_slow(
+ #[inline(never)]
+ extern "C-unwind" fn call_structural_visit_hook_slow_raw(
visitor: StructuralVisitorHandle,
- raw: TVMFFIAny,
- attr: TVMFFIAny,
- ) -> Any {
+ value_type_index: i32,
+ value_small_str_len: u32,
+ value_data: TVMFFIAnyDataUnion,
+ attr_type_index: i32,
+ attr_small_str_len: u32,
+ attr_data: TVMFFIAnyDataUnion,
+ ) -> TVMFFIAny {
+ // Construct addressable packed-call arguments only on this cold path.
+ let raw = TVMFFIAny {
+ type_index: value_type_index,
+ small_str_len: value_small_str_len,
+ data_union: value_data,
+ };
+ let attr = TVMFFIAny {
+ type_index: attr_type_index,
+ small_str_len: attr_small_str_len,
+ data_union: attr_data,
+ };
// Packed hooks (including Python class methods) can return any
well-formed
// Any through safe APIs. Their result tag needs validation before
decoding.
let result = unsafe {
@@ -488,12 +593,14 @@ impl StructuralVisitEngine {
])
})
};
- match result {
+ let result = match result {
Ok(value) => checked_hook_result(value),
Err(error) => error.into(),
- }
+ };
+ unsafe { Any::into_raw_ffi_any(result) }
}
+ #[inline(always)]
unsafe fn borrowed_visitor_view<'a>(visitor: StructuralVisitorHandle) ->
AnyView<'a> {
let object = visitor.cast::<TVMFFIObject>();
let mut raw = TVMFFIAny::new();
@@ -521,6 +628,7 @@ impl StructuralVisitEngine {
previous: i32,
}
impl Drop for Restore {
+ #[inline(always)]
fn drop(&mut self) {
unsafe { (*self.visitor).def_region_mode = self.previous };
}
@@ -531,6 +639,7 @@ impl StructuralVisitEngine {
}
}
+#[inline(always)]
fn dispatcher_context<D>(dispatch: &mut D) -> *mut c_void {
std::ptr::from_mut(dispatch).cast()
}
@@ -542,15 +651,22 @@ fn wrong_visitor_dispatcher() -> ! {
}
#[inline(always)]
-fn visit_result(result: Result<Option<VisitInterrupt>>, raw: TVMFFIAny) -> Any
{
+fn visit_result_raw(result: Result<Option<VisitInterrupt>>, raw: TVMFFIAny) ->
TVMFFIAny {
match result {
- Ok(None) => Any::new(),
- Ok(Some(interrupt)) => interrupt.into(),
- Err(error) => with_value_context(error.into(), raw),
+ Ok(None) => TVMFFIAny::new(),
+ Ok(Some(interrupt)) => unsafe {
Any::into_raw_ffi_any(interrupt.into()) },
+ Err(error) => {
+ core::hint::cold_path();
+ StructuralVisitEngine::with_visit_context_raw(
+ unsafe { Any::into_raw_ffi_any(error.into()) },
+ raw,
+ )
+ }
}
}
// A packed hook's owning Any is valid, but its result type is dynamic.
+#[inline(always)]
fn checked_hook_result(value: Any) -> Any {
let tag = value.type_index();
if tag == TVMFFITypeIndex::kTVMFFINone as i32
@@ -569,31 +685,14 @@ fn checked_hook_result(value: Any) -> Any {
}
/// Consume exactly one Error reference after its integer tag was checked.
+#[inline(always)]
unsafe fn take_error(value: Any) -> Error {
use crate::type_traits::AnyCompatible;
let mut raw = Any::into_raw_ffi_any(value);
Error::move_from_any_after_check(&mut raw)
}
-#[inline]
-fn with_value_context(result: Any, value: TVMFFIAny) -> Any {
- if result.type_index() != TVMFFITypeIndex::kTVMFFIError as i32
- || value.type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32
- {
- result
- } else {
- // SAFETY: the Error tag establishes the owned object's representation.
- let error = with_visit_error_context(unsafe { take_error(result) },
value);
- with_structural_error_context(
- error,
- "visit",
- &format!("object `{}`", type_key_of(value.type_index)),
- )
- .into()
- }
-}
-
-#[inline]
+#[inline(always)]
pub(crate) fn field_def_region(field: &TVMFFIFieldInfo, inherited:
DefRegionKind) -> DefRegionKind {
// Precedence: a pattern region propagates; entering any kind inside it
has no effect.
if inherited == DefRegionKind::Pattern {
@@ -611,7 +710,7 @@ pub(crate) fn field_def_region(field: &TVMFFIFieldInfo,
inherited: DefRegionKind
/// the FreeVar's own reflected children: nested free vars there must resolve
/// against an outer binding instead of rebinding. Mirrors C++
/// `VisitReflectedFieldsExpected`.
-#[inline]
+#[inline(always)]
pub(crate) fn free_var_child_region(
inherited: DefRegionKind,
structural_eq_hash_kind: i32,
@@ -625,9 +724,10 @@ pub(crate) fn free_var_child_region(
}
}
-#[inline]
+#[inline(always)]
fn with_error_context(result: Any, frame: &str) -> Any {
if result.type_index() == TVMFFITypeIndex::kTVMFFIError as i32 {
+ core::hint::cold_path();
// SAFETY: direct tag check establishes the Error representation.
with_structural_error_context(unsafe { take_error(result) }, "visit",
frame).into()
} else {
@@ -647,6 +747,22 @@ static STRUCTURAL_VISIT_COLUMN: LazyLock<TypeAttrColumn> =
LazyLock::new(|| {
TypeAttrColumn::new(STRUCTURAL_VISIT_ATTR).expect("missing __s_visit__
column")
});
+#[inline(always)]
+fn structural_visit_column() -> TypeAttrColumn {
+ if let Some(column) = LazyLock::get(&STRUCTURAL_VISIT_COLUMN) {
+ *column
+ } else {
+ initialize_structural_visit_column()
+ }
+}
+
+// Keep Once's initialization descriptors out of each recursive caller's frame.
+#[cold]
+#[inline(never)]
+fn initialize_structural_visit_column() -> TypeAttrColumn {
+ *LazyLock::force(&STRUCTURAL_VISIT_COLUMN)
+}
+
pub(crate) fn type_key_of(type_index: i32) -> String {
unsafe {
let info = TVMFFIGetTypeInfo(type_index);
@@ -738,11 +854,11 @@ impl<D> StructuralVisitorView<'_, D> {
// SAFETY: the view ties this runtime to D; the check establishes the
// instance, and publishing uses the caller's actual exclusive
reborrow.
unsafe {
- finish(StructuralVisitEngine::visit(
+
visit_result_from_any(Any::from_raw_ffi_any(StructuralVisitEngine::visit_raw(
self.runtime,
dispatch,
raw_of(value.into()),
- ))
+ )))
}
}
/// Run default descent through registered hooks or reflected fields.
@@ -760,10 +876,12 @@ impl<D> StructuralVisitorView<'_, D> {
// the supplied borrow must identify this runtime's dispatcher
instance.
unsafe {
let context =
StructuralVisitEngine::checked_dispatcher(self.runtime, dispatch);
- finish(StructuralVisitEngine::default_visit_raw(
- self.runtime,
- context,
- raw_of(value.into()),
+ visit_result_from_any(Any::from_raw_ffi_any(
+ StructuralVisitEngine::default_visit_raw(
+ self.runtime,
+ context,
+ raw_of(value.into()),
+ ),
))
}
}
@@ -782,8 +900,14 @@ where
Option<Result<Option<VisitInterrupt>>>,
>,
{
+ #[inline(always)]
fn visit_root(mut self, root: AnyView<'_>) ->
Result<Option<VisitInterrupt>> {
- finish(StructuralVisitEngine::run_visitor(raw_of(root), &mut self))
+ visit_result_from_any(unsafe {
+ Any::from_raw_ffi_any(StructuralVisitEngine::run_visitor_raw(
+ raw_of(root),
+ &mut self,
+ ))
+ })
}
}
@@ -794,8 +918,11 @@ where
Option<Result<Option<VisitInterrupt>>>,
>,
{
+ #[inline(always)]
fn visit_root(self, root: AnyView<'_>) -> Result<Option<VisitInterrupt>> {
- finish(StructuralVisitEngine::run_visitor(raw_of(root), self))
+ visit_result_from_any(unsafe {
+
Any::from_raw_ffi_any(StructuralVisitEngine::run_visitor_raw(raw_of(root),
self))
+ })
}
}
@@ -818,11 +945,14 @@ where
visitor.visit_root(root.into())
}
-pub(super) fn finish(result: Any) -> Result<Option<VisitInterrupt>> {
+#[inline(always)]
+pub(super) fn visit_result_from_any(result: Any) ->
Result<Option<VisitInterrupt>> {
let tag = result.type_index();
if tag == TVMFFITypeIndex::kTVMFFINone as i32 {
- Ok(None)
- } else if tag == TVMFFITypeIndex::kTVMFFIError as i32 {
+ return Ok(None);
+ }
+ core::hint::cold_path();
+ if tag == TVMFFITypeIndex::kTVMFFIError as i32 {
// SAFETY: internal results obey the documented owning carrier
protocol.
Err(unsafe { take_error(result) })
} else {
diff --git a/rust/tvm-ffi/src/extra/structural_walk.rs
b/rust/tvm-ffi/src/extra/structural_walk.rs
index 61e04a4d..8a17b280 100644
--- a/rust/tvm-ffi/src/extra/structural_walk.rs
+++ b/rust/tvm-ffi/src/extra/structural_walk.rs
@@ -24,7 +24,8 @@
use super::structural_common::{impl_callback_chain_tuple_arities, raw_of};
use super::structural_visit::{
- finish, DefRegionKind, StructuralVisitEngine, StructuralVisitorView,
VisitInterrupt,
+ visit_result_from_any, DefRegionKind, StructuralVisitEngine,
StructuralVisitorView,
+ VisitInterrupt,
};
use crate::{Any, AnyView, Dispatcher, ObjectCore, Result};
use std::marker::PhantomData;
@@ -47,6 +48,7 @@ pub enum WalkResult {
impl WalkResult {
/// Halt traversal with an FFI-compatible payload.
+ #[inline(always)]
pub fn interrupt_with<T: Into<Any>>(payload: T) -> Self {
Self::InterruptWith(payload.into())
}
@@ -61,12 +63,14 @@ pub trait IntoWalkResult {
}
impl IntoWalkResult for WalkResult {
+ #[inline(always)]
fn into_walk_result(self) -> Result<WalkResult> {
Ok(self)
}
}
impl IntoWalkResult for Result<WalkResult> {
+ #[inline(always)]
fn into_walk_result(self) -> Result<WalkResult> {
self
}
@@ -134,6 +138,7 @@ pub enum ByWalkDispatch {}
impl<'a, V: WalkDispatch> IntoWalker<ByWalkDispatch> for &'a mut V {
type Walker = DispatchWalker<&'a mut V>;
+ #[inline(always)]
fn into_walker(self) -> Self::Walker {
DispatchWalker { walker: self }
}
@@ -148,6 +153,7 @@ pub struct DispatchWalker<V> {
impl<V> NativeContext for DispatchWalker<V> {}
impl<V: WalkDispatch> NativeVisit for DispatchWalker<&mut V> {
+ #[inline(always)]
fn visit(&mut self, value: AnyView<'_>, def_region_kind: DefRegionKind) ->
Result<WalkResult> {
self.walker
.dispatch_walk(value, def_region_kind)
@@ -182,6 +188,7 @@ where
O: IntoWalkResult,
{
type Walker = ChainWalker<F, ByCatchAllLink>;
+ #[inline(always)]
fn into_walker(self) -> Self::Walker {
ChainWalker::new(self)
}
@@ -193,6 +200,7 @@ where
O: IntoWalkResult,
{
type Walker = ChainWalker<F, ByCatchAllKindLink>;
+ #[inline(always)]
fn into_walker(self) -> Self::Walker {
ChainWalker::new(self)
}
@@ -272,7 +280,7 @@ where
T: crate::type_traits::AnyCompatible,
O: IntoWalkResult,
{
- #[inline]
+ #[inline(always)]
fn try_call(
&mut self,
value: AnyView<'_>,
@@ -293,7 +301,7 @@ where
T: crate::type_traits::AnyCompatible,
O: IntoWalkResult,
{
- #[inline]
+ #[inline(always)]
fn try_call(
&mut self,
value: AnyView<'_>,
@@ -314,7 +322,7 @@ where
N: ObjectCore,
O: IntoWalkResult,
{
- #[inline]
+ #[inline(always)]
fn try_call(
&mut self,
value: AnyView<'_>,
@@ -335,7 +343,7 @@ where
N: ObjectCore,
O: IntoWalkResult,
{
- #[inline]
+ #[inline(always)]
fn try_call(
&mut self,
value: AnyView<'_>,
@@ -355,7 +363,7 @@ where
F: for<'a> FnMut(AnyView<'a>) -> O,
O: IntoWalkResult,
{
- #[inline]
+ #[inline(always)]
fn try_call(
&mut self,
value: AnyView<'_>,
@@ -373,7 +381,7 @@ where
F: for<'a> FnMut(AnyView<'a>, DefRegionKind) -> O,
O: IntoWalkResult,
{
- #[inline]
+ #[inline(always)]
fn try_call(
&mut self,
value: AnyView<'_>,
@@ -390,7 +398,7 @@ pub struct ByChainLink<Markers>(PhantomData<fn(Markers)>);
pub enum ByWalkDispatchLink {}
impl<V: WalkDispatch> WalkChainLink<ByWalkDispatchLink> for &mut V {
- #[inline]
+ #[inline(always)]
fn try_call(
&mut self,
value: AnyView<'_>,
@@ -408,7 +416,7 @@ pub struct ChainWalker<Link, Marker> {
}
impl<Link, Marker> ChainWalker<Link, Marker> {
- #[inline]
+ #[inline(always)]
fn new(link: Link) -> Self {
ChainWalker {
link,
@@ -423,7 +431,7 @@ impl<Link, Marker> NativeVisit for ChainWalker<Link, Marker>
where
Link: WalkChainLink<Marker>,
{
- #[inline]
+ #[inline(always)]
fn visit(&mut self, value: AnyView<'_>, def_region_kind: DefRegionKind) ->
Result<WalkResult> {
self.link
.try_call(value, def_region_kind)
@@ -443,7 +451,7 @@ macro_rules! impl_chain_link {
where
$($F: WalkChainLink<$M>,)+
{
- #[inline]
+ #[inline(always)]
fn try_call(
&mut self,
value: AnyView<'_>,
@@ -463,6 +471,7 @@ macro_rules! impl_chain_link {
$($F: WalkChainLink<$M>,)+
{
type Walker = ChainWalker<($($F,)+), ByChainLink<($($M,)+)>>;
+ #[inline(always)]
fn into_walker(self) -> Self::Walker {
ChainWalker::new(self)
}
@@ -482,6 +491,7 @@ macro_rules! impl_bare_link_walker {
O: IntoWalkResult,
{
type Walker = ChainWalker<F, $marker<T>>;
+ #[inline(always)]
fn into_walker(self) -> Self::Walker {
ChainWalker::new(self)
}
@@ -500,8 +510,11 @@ impl_bare_link_walker!(
/// Internal callback protocol used by [`IntoWalker`].
#[doc(hidden)]
pub trait NativeVisit: Sized + sealed::NativeContext {
+ #[inline(always)]
fn walk_root(self, root: AnyView<'_>, order: WalkOrder) ->
Result<Option<VisitInterrupt>> {
- finish(run_walker(raw_of(root), self, order))
+ visit_result_from_any(unsafe {
+ Any::from_raw_ffi_any(run_walker_raw(raw_of(root), self, order))
+ })
}
fn visit(&mut self, value: AnyView<'_>, def_region_kind: DefRegionKind) ->
Result<WalkResult>;
@@ -510,6 +523,7 @@ pub trait NativeVisit: Sized + sealed::NativeContext {
impl<V: NativeVisit> NativeContext for &mut V {}
impl<V: NativeVisit> NativeVisit for &mut V {
+ #[inline(always)]
fn visit(&mut self, value: AnyView<'_>, def_region_kind: DefRegionKind) ->
Result<WalkResult> {
(**self).visit(value, def_region_kind)
}
@@ -557,31 +571,56 @@ impl<V: NativeVisit, const PRE_ORDER: bool>
WalkVisitor<V, PRE_ORDER> {
let kind = visitor.def_region_kind();
if !PRE_ORDER {
if let Some(interrupt) = visitor.default_visit(self, value)? {
+ core::hint::cold_path();
return Ok(Some(interrupt));
}
}
match self.observer.visit(value, kind)? {
WalkResult::Advance if PRE_ORDER => visitor.default_visit(self,
value),
- WalkResult::Advance | WalkResult::Skip => Ok(None),
- WalkResult::Interrupt =>
Ok(Some(VisitInterrupt::with(Any::new()))),
- WalkResult::InterruptWith(payload) =>
Ok(Some(VisitInterrupt::with(payload))),
+ WalkResult::Advance => Ok(None),
+ WalkResult::Skip => {
+ core::hint::cold_path();
+ Ok(None)
+ }
+ WalkResult::Interrupt => {
+ core::hint::cold_path();
+ Ok(Some(VisitInterrupt::with(Any::new())))
+ }
+ WalkResult::InterruptWith(payload) => {
+ core::hint::cold_path();
+ Ok(Some(VisitInterrupt::with(payload)))
+ }
}
}
}
-fn run_walker<D: NativeVisit>(
+#[inline(always)]
+fn run_walker_raw<D: NativeVisit>(
root: crate::tvm_ffi_sys::TVMFFIAny,
observer: D,
order: WalkOrder,
-) -> Any {
- match order {
+) -> crate::tvm_ffi_sys::TVMFFIAny {
+ // Own the result while dropping the observer: user capture destructors can
+ // unwind after the native callback has returned.
+ let result = match order {
WalkOrder::PreOrder => {
- StructuralVisitEngine::run_visitor(root, &mut WalkVisitor::<D,
true> { observer })
+ let mut visitor = WalkVisitor::<D, true> { observer };
+ let result = unsafe {
+
Any::from_raw_ffi_any(StructuralVisitEngine::run_visitor_raw(root, &mut
visitor))
+ };
+ drop(visitor);
+ result
}
WalkOrder::PostOrder => {
- StructuralVisitEngine::run_visitor(root, &mut WalkVisitor::<D,
false> { observer })
+ let mut visitor = WalkVisitor::<D, false> { observer };
+ let result = unsafe {
+
Any::from_raw_ffi_any(StructuralVisitEngine::run_visitor_raw(root, &mut
visitor))
+ };
+ drop(visitor);
+ result
}
- }
+ };
+ unsafe { Any::into_raw_ffi_any(result) }
}
/// Walk `root` with an observer, the Rust analog of C++
diff --git a/rust/tvm-ffi/src/object.rs b/rust/tvm-ffi/src/object.rs
index e0b9137c..4828f827 100644
--- a/rust/tvm-ffi/src/object.rs
+++ b/rust/tvm-ffi/src/object.rs
@@ -625,7 +625,7 @@ impl<T: ObjectCore> ObjectArc<T> {
/// `ptr` must own one strong reference to a live, stable managed
allocation
/// whose dynamic type is compatible with `T` and whose payload is valid.
/// Ownership of that reference is transferred into the returned handle.
- #[inline]
+ #[inline(always)]
pub unsafe fn from_raw(ptr: *const T) -> Self {
Self {
ptr: std::ptr::NonNull::new_unchecked(ptr as *mut T),
@@ -640,7 +640,7 @@ impl<T: ObjectCore> ObjectArc<T> {
///
/// # Returns
/// * `*const T` - The raw pointer
- #[inline]
+ #[inline(always)]
pub unsafe fn into_raw(this: Self) -> *const T {
let droped_this = std::mem::ManuallyDrop::new(this);
droped_this.ptr.as_ptr() as *const T
@@ -719,6 +719,7 @@ impl<T: ObjectCore> Deref for ObjectArc<T> {
// implement Drop for ObjectArc
impl<T: ObjectCore> Drop for ObjectArc<T> {
+ #[inline(always)]
fn drop(&mut self) {
unsafe { unsafe_::dec_ref(self.ptr.as_ptr().cast::<TVMFFIObject>()) }
}
diff --git a/rust/tvm-ffi/src/reflection.rs b/rust/tvm-ffi/src/reflection.rs
index f0b2e31c..019573d5 100644
--- a/rust/tvm-ffi/src/reflection.rs
+++ b/rust/tvm-ffi/src/reflection.rs
@@ -69,6 +69,7 @@ impl TypeAttrColumn {
}
/// Copy one borrowed cell without taking ownership.
+ #[inline(always)]
pub(crate) fn get_raw(self, type_index: i32) -> Option<TVMFFIAny> {
unsafe {
let column = self.0.as_ref();
