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-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 31f6c2b9d [Rust] Borrow the arguments of read-only query APIs through
IntoObjectRef (#206)
31f6c2b9d is described below
commit 31f6c2b9d54f416040d82c0783714d5a4bab7188
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Nov 5 19:09:31 2026 -0500
[Rust] Borrow the arguments of read-only query APIs through IntoObjectRef
(#206)
Read-only queries took their object arguments as `impl Into<T>`, so a
caller passing `&x` paid a reference-count bump for a value the call
only reads. They now take `impl IntoObjectRef<T>` and pass the borrow,
as `IntoExpr` does for expressions: `&x` and its subtypes cost nothing,
and owned values still work.
## Changes
- `IRModule::lookup`, `contains_global_var` and `remove_global_var`
borrow the `GlobalVar`.
- `Analyzer::const_int_bound_is_bound` and `tensor_is_scalar` borrow the
`Var`.
- `ThreadScope::is_higher_than` borrows `other`.
- `analysis::verify_memory`, `check_well_formed` and
`check_well_formed_module` borrow the function or module.
- The in-place mutation check in
`StructuralMutatorWithConstraintContext` borrows the object with
`downcast_ref` instead of copying it with `try_as`, so the uniqueness
test is a strong count of 1.
APIs that hand ownership on (`Pass::run`, `build`,
`IRModule::update_func`) stay by value, so a unique reference can be
modified in place.
---
rust/src/ir/module.rs | 14 ++++++++------
rust/src/sym/analyzer.rs | 5 +++--
rust/src/tirx/analysis/analysis.rs | 16 ++++++++++------
.../tirx/structural_mutate_with_constraint_context.rs | 11 ++++-------
rust/src/tirx/tensor.rs | 6 ++++--
rust/src/tirx/thread_scope.rs | 6 ++++--
6 files changed, 33 insertions(+), 25 deletions(-)
diff --git a/rust/src/ir/module.rs b/rust/src/ir/module.rs
index 058b9f988..a8c2f0f73 100644
--- a/rust/src/ir/module.rs
+++ b/rust/src/ir/module.rs
@@ -65,6 +65,7 @@ use std::borrow::Borrow;
use std::ops::Deref;
use tvm_ffi::Any;
use tvm_ffi::Array;
+use tvm_ffi::IntoObjectRef;
use tvm_ffi::Map;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
@@ -217,8 +218,9 @@ impl IRModule {
}
/// Remove the function of `var`, if the module has one
(`ir.IRModuleRemove`).
- pub fn remove_global_var(&self, var: impl Into<GlobalVar>) -> Result<()> {
- tvm_ffi::cached_global_func!("ir.IRModuleRemove").call_tuple((self,
&var.into()))?;
+ pub fn remove_global_var(&self, var: impl IntoObjectRef<GlobalVar>) ->
Result<()> {
+ tvm_ffi::cached_global_func!("ir.IRModuleRemove")
+ .call_tuple((self, var.into_object_ref().borrow()))?;
Ok(())
}
@@ -230,9 +232,9 @@ impl IRModule {
}
/// Return whether the module defines a function for `var`
(`ir.IRModuleContains`).
- pub fn contains_global_var(&self, var: impl Into<GlobalVar>) ->
Result<bool> {
+ pub fn contains_global_var(&self, var: impl IntoObjectRef<GlobalVar>) ->
Result<bool> {
tvm_ffi::cached_global_func!("ir.IRModuleContains")
- .call_tuple((self, &var.into()))?
+ .call_tuple((self, var.into_object_ref().borrow()))?
.try_into()
}
@@ -265,9 +267,9 @@ impl IRModule {
/// # Errors
///
/// Returns an `InternalError` if the module has no function for `var`.
- pub fn lookup(&self, var: impl Into<GlobalVar>) -> Result<BaseFunc> {
+ pub fn lookup(&self, var: impl IntoObjectRef<GlobalVar>) ->
Result<BaseFunc> {
tvm_ffi::cached_global_func!("ir.IRModuleLookup")
- .call_tuple((self, &var.into()))?
+ .call_tuple((self, var.into_object_ref().borrow()))?
.try_into()
}
diff --git a/rust/src/sym/analyzer.rs b/rust/src/sym/analyzer.rs
index 9163ede65..7f19af5df 100644
--- a/rust/src/sym/analyzer.rs
+++ b/rust/src/sym/analyzer.rs
@@ -83,6 +83,7 @@ use crate::sym::int_set::IntSet;
use std::borrow::Borrow;
use std::ops::Deref;
use tvm_ffi::Error;
+use tvm_ffi::IntoObjectRef;
use tvm_ffi::Map;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
@@ -178,9 +179,9 @@ impl Analyzer {
}
/// Return whether `var` has constant bounds recorded
(`sym.AnalyzerConstIntBoundIsBound`).
- pub fn const_int_bound_is_bound(&self, var: impl Into<Var>) ->
Result<bool> {
+ pub fn const_int_bound_is_bound(&self, var: impl IntoObjectRef<Var>) ->
Result<bool> {
tvm_ffi::cached_global_func!("sym.AnalyzerConstIntBoundIsBound")
- .call_tuple((self, &var.into()))?
+ .call_tuple((self, var.into_object_ref().borrow()))?
.try_into()
}
diff --git a/rust/src/tirx/analysis/analysis.rs
b/rust/src/tirx/analysis/analysis.rs
index f29a6d567..656919db4 100644
--- a/rust/src/tirx/analysis/analysis.rs
+++ b/rust/src/tirx/analysis/analysis.rs
@@ -57,6 +57,7 @@ use crate::tirx::function::Function;
use std::borrow::Borrow;
use tvm_ffi::Any;
use tvm_ffi::Array;
+use tvm_ffi::IntoObjectRef;
use tvm_ffi::Result;
// tvm-ffi-stubgen(end)
@@ -65,9 +66,9 @@ use tvm_ffi::Result;
///
/// A function without a target attribute, or not of the default calling
convention, is not checked
/// and passes.
-pub fn verify_memory(func: impl Into<Function>) -> Result<bool> {
+pub fn verify_memory(func: impl IntoObjectRef<Function>) -> Result<bool> {
tvm_ffi::cached_global_func!("tirx.analysis.VerifyMemory")
- .call_tuple((&func.into(),))?
+ .call_tuple((func.into_object_ref().borrow(),))?
.try_into()
}
@@ -85,17 +86,20 @@ pub fn undefined_vars(
/// Return whether `func` is well formed; invalid IR raises an error instead
if `assert_mode`
/// (`tirx.analysis.CheckWellFormed`).
-pub fn check_well_formed(func: impl Into<Function>, assert_mode: bool) ->
Result<bool> {
+pub fn check_well_formed(func: impl IntoObjectRef<Function>, assert_mode:
bool) -> Result<bool> {
tvm_ffi::cached_global_func!("tirx.analysis.CheckWellFormed")
- .call_tuple((&func.into(), assert_mode))?
+ .call_tuple((func.into_object_ref().borrow(), assert_mode))?
.try_into()
}
/// Return whether the TIRx functions of `module` are well formed; invalid IR
raises an error
/// instead if `assert_mode` (`tirx.analysis.CheckWellFormed`).
-pub fn check_well_formed_module(module: impl Into<IRModule>, assert_mode:
bool) -> Result<bool> {
+pub fn check_well_formed_module(
+ module: impl IntoObjectRef<IRModule>,
+ assert_mode: bool,
+) -> Result<bool> {
tvm_ffi::cached_global_func!("tirx.analysis.CheckWellFormed")
- .call_tuple((&module.into(), assert_mode))?
+ .call_tuple((module.into_object_ref().borrow(), assert_mode))?
.try_into()
}
diff --git a/rust/src/tirx/structural_mutate_with_constraint_context.rs
b/rust/src/tirx/structural_mutate_with_constraint_context.rs
index 471848d92..758579e94 100644
--- a/rust/src/tirx/structural_mutate_with_constraint_context.rs
+++ b/rust/src/tirx/structural_mutate_with_constraint_context.rs
@@ -95,13 +95,10 @@ unsafe impl StructuralMutatorBase for
StructuralMutatorWithConstraintContextBase
let hook = if value.inplace_mode() == InplaceMode::Allow {
lookup(*INPLACE).filter(|_| {
// Permission alone is insufficient after a callback retains an
- // alias. The supported owning count includes this temporary;
- // release it before invoking a hook or entering native
descent.
- borrowed.try_as::<ObjectRef>().is_some_and(|owner| {
- let unique =
ObjectArc::strong_count(ObjectRef::data(&owner)) == 2;
- drop(owner);
- unique
- })
+ // alias. Borrow in place so the count is the slot's own
reference.
+ borrowed
+ .downcast_ref::<ObjectRef>()
+ .is_some_and(|owner|
ObjectArc::strong_count(ObjectRef::data(owner)) == 1)
})
} else {
None
diff --git a/rust/src/tirx/tensor.rs b/rust/src/tirx/tensor.rs
index 7679d5537..720b8c43e 100644
--- a/rust/src/tirx/tensor.rs
+++ b/rust/src/tirx/tensor.rs
@@ -76,9 +76,11 @@ use crate::ir::support;
use crate::runtime::MemoryScope;
use crate::sym::analyzer::Analyzer;
use crate::tirx::layout::Layout;
+use std::borrow::Borrow;
use std::ops::Deref;
use tvm_ffi::AnyView;
use tvm_ffi::Array;
+use tvm_ffi::IntoObjectRef;
use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
@@ -228,9 +230,9 @@ impl TensorType {
/// # Errors
///
/// Returns an error if the type of `tensor` is not a [`TensorType`].
-pub fn tensor_is_scalar(tensor: impl Into<Var>, alloc_or_decl: bool) ->
Result<bool> {
+pub fn tensor_is_scalar(tensor: impl IntoObjectRef<Var>, alloc_or_decl: bool)
-> Result<bool> {
tvm_ffi::cached_global_func!("tirx.TensorIsScalar")
- .call_tuple((&tensor.into(), alloc_or_decl))?
+ .call_tuple((tensor.into_object_ref().borrow(), alloc_or_decl))?
.try_into()
}
diff --git a/rust/src/tirx/thread_scope.rs b/rust/src/tirx/thread_scope.rs
index 9bc9fd264..9c6956d82 100644
--- a/rust/src/tirx/thread_scope.rs
+++ b/rust/src/tirx/thread_scope.rs
@@ -41,7 +41,9 @@
// tvm-ffi-stubgen(import-object): tvm_ffi::String
// tvm-ffi-stubgen(begin): import-section
+use std::borrow::Borrow;
use std::ops::Deref;
+use tvm_ffi::IntoObjectRef;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
use tvm_ffi::Result;
@@ -116,9 +118,9 @@ impl ThreadScope {
/// # Errors
///
/// Returns an error if the two scopes are of different backends.
- pub fn is_higher_than(&self, other: impl Into<ThreadScope>) ->
Result<bool> {
+ pub fn is_higher_than(&self, other: impl IntoObjectRef<ThreadScope>) ->
Result<bool> {
tvm_ffi::cached_global_func!("tirx.ScopeHigher")
- .call_tuple((self, &other.into()))?
+ .call_tuple((self, other.into_object_ref().borrow()))?
.try_into()
}
}
