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 a6bb154d7 [Rust] Bump tvm-ffi to f0520693; cache generated field
getters (#227)
a6bb154d7 is described below
commit a6bb154d723e9f72184faa563fce9d17b1b02e3b
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Nov 7 19:28:48 2026 -0500
[Rust] Bump tvm-ffi to f0520693; cache generated field getters (#227)
Bump 3rdparty/tvm-ffi from c244bc5f to apache/tvm-ffi f0520693
(apache/tvm-ffi#961), whose Rust stub generator caches each field
accessor's `FieldGetter` in a `ColdInitLazyLock` instead of resolving it
by name on every call.
- Regenerated bindings: the 20 remaining generated field accessors use
the cached getter.
- Target: drop the mutable `str_repr_` cache; `str()` serializes the
config on each call.
- Register empty reflection for the field-less `tirx.Layout`,
`sym.IterMapExpr`, and `sym.IntSet` bases so their layout is known.
- `SplitHostDevice` and `SpecializeDeviceHelpers` compare targets
structurally instead of by their JSON string.
- Hide the canonical simplifier's internal `sym.CanonicalExpr`,
`sym.SplitExpr`, and `sym.SumExpr` from the Rust bindings.
- Regenerated bindings: `Target`, `TileLayout`, `ComposeLayout`,
`IterSplitExpr`, `IterSumExpr`, `IntervalSet`, and their bases are now
plain `repr(C)` field mirrors.
- Drop the field-wise allocators of these constructed types
(`from_fields` on the six complete types, `new` on the `Layout`,
`IntSet`, and `IterMapExpr` bases and on their `*Obj` types): their C++
constructors parse configs, derive masks, and check invariants, which a
field-wise allocator would bypass. The stub generator has no directive
to omit an allocator, so these nine bindings are hand-written (`skip` +
`ty-map`) with the generated layout, field reads, and upcasts.
---
3rdparty/tvm-ffi | 2 +-
include/tvm/target/target.h | 6 --
rust/src/ir/expr.rs | 6 +-
rust/src/ir/transform/core.rs | 17 +++-
rust/src/lib.rs | 4 +-
rust/src/sym.rs | 10 ++
rust/src/sym/analyzer.rs | 114 +---------------------
rust/src/sym/int_set.rs | 45 ++++-----
rust/src/sym/iter_affine_map.rs | 85 +++++++---------
rust/src/target.rs | 58 +++++------
rust/src/tirx.rs | 4 +
rust/src/tirx/layout.rs | 97 ++++++++----------
rust/src/tirx/structural.rs | 11 ++-
rust/src/tirx/tile/dispatch.rs | 48 +++++++--
rust/src/tirx/transform/function_pass.rs | 5 +-
rust/tests/sym/analyzer.rs | 10 +-
rust/tests/sym/int_set.rs | 6 +-
rust/tests/sym/iter_affine_map.rs | 18 ++--
rust/tests/target/main.rs | 31 +++---
rust/tests/tirx/layout.rs | 2 +-
src/backend/cuda/transform/bind_backend_config.cc | 11 ++-
src/sym/int_set.cc | 5 +-
src/sym/iter_affine_map.cc | 1 +
src/target/target.cc | 8 +-
src/tirx/ir/layout/tile_core.cc | 5 +-
src/tirx/transform/split_host_device.cc | 3 +-
26 files changed, 260 insertions(+), 352 deletions(-)
diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi
index c244bc5fc..f05206931 160000
--- a/3rdparty/tvm-ffi
+++ b/3rdparty/tvm-ffi
@@ -1 +1 @@
-Subproject commit c244bc5fcbaac54bbc6f0563fccf302cb570388a
+Subproject commit f052069318165cb60f3fc7dd0d3252db9eddffee
diff --git a/include/tvm/target/target.h b/include/tvm/target/target.h
index b6d4cc343..4bf23a7f5 100644
--- a/include/tvm/target/target.h
+++ b/include/tvm/target/target.h
@@ -147,12 +147,6 @@ class TargetObj : public ffi::Object {
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("target.Target", TargetObj, ffi::Object);
// Shared instances are mutated in place, so the class is mutable.
static constexpr bool _type_mutable = true;
-
- private:
- /*! \brief Internal string repr. */
- mutable ffi::String str_repr_;
-
- friend class TargetInternal;
};
/*!
diff --git a/rust/src/ir/expr.rs b/rust/src/ir/expr.rs
index d1eb86355..d2878f6ac 100644
--- a/rust/src/ir/expr.rs
+++ b/rust/src/ir/expr.rs
@@ -2315,13 +2315,11 @@ crate::def_expr_op_over_loads!(crate::ir::prim::Ramp);
crate::def_expr_op_over_loads!(crate::ir::prim::Broadcast);
crate::def_expr_op_over_loads!(crate::ir::prim::Shuffle);
-// These symbolic reference classes also inherit C++ PrimExpr. The opaque Rust
mirrors retain
-// their existing allocation restrictions; attaching operators only forwards
their handles.
+// These symbolic reference classes also inherit C++ PrimExpr; attaching
operators only forwards
+// their handles.
crate::def_expr_op_over_loads!(crate::sym::iter_affine_map::IterMapExpr);
crate::def_expr_op_over_loads!(crate::sym::iter_affine_map::IterSplitExpr);
crate::def_expr_op_over_loads!(crate::sym::iter_affine_map::IterSumExpr);
-crate::def_expr_op_over_loads!(crate::sym::analyzer::SplitExpr);
-crate::def_expr_op_over_loads!(crate::sym::analyzer::SumExpr);
//
---------------------------------------------------------------------------------------------
// What the constant constructors check
diff --git a/rust/src/ir/transform/core.rs b/rust/src/ir/transform/core.rs
index 4d7b63788..0a8fa6322 100644
--- a/rust/src/ir/transform/core.rs
+++ b/rust/src/ir/transform/core.rs
@@ -48,6 +48,7 @@ use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
use tvm_ffi::String;
+use tvm_ffi::extra::cold_init_lazy_lock::ColdInitLazyLock;
use tvm_ffi::reflection::FieldGetter;
// tvm-ffi-stubgen(end)
@@ -198,7 +199,9 @@ impl Deref for ModulePassObj {
#[allow(missing_docs)]
impl ModulePassObj {
pub fn pass_info(&self) -> Result<PassInfo> {
- FieldGetter::new(Self::type_index(), "pass_info")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(ModulePassObj::type_index(), "pass_info"));
+ GETTER.clone()?.get(self)
}
}
@@ -260,11 +263,15 @@ impl Deref for SequentialObj {
#[allow(missing_docs)]
impl SequentialObj {
pub fn pass_info(&self) -> Result<PassInfo> {
- FieldGetter::new(Self::type_index(), "pass_info")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(SequentialObj::type_index(), "pass_info"));
+ GETTER.clone()?.get(self)
}
pub fn passes(&self) -> Result<Array<Pass>> {
- FieldGetter::new(Self::type_index(), "passes")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(SequentialObj::type_index(), "passes"));
+ GETTER.clone()?.get(self)
}
}
@@ -502,7 +509,9 @@ pub struct PassInstrumentObj {
#[allow(missing_docs)]
impl PassInstrumentObj {
pub fn name(&self) -> Result<String> {
- FieldGetter::new(Self::type_index(), "name")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(PassInstrumentObj::type_index(), "name"));
+ GETTER.clone()?.get(self)
}
}
// tvm-ffi-stubgen(end)
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index ed14c7cd9..bcece86e9 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -48,10 +48,10 @@
//! use tvm_tirx::tvm_ffi::{Function, String};
//!
//! let target = Target::new("llvm")?;
-//! assert_eq!(target.kind()?.name()?, "llvm");
+//! assert_eq!(target.kind.name()?, "llvm");
//!
//! let same =
Function::get_global("target.Target")?.call_tuple((String::from("llvm"),))?;
-//! assert_eq!(Target::try_from(same)?.kind()?.name()?, "llvm");
+//! assert_eq!(Target::try_from(same)?.kind.name()?, "llvm");
//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
//! ```
//!
diff --git a/rust/src/sym.rs b/rust/src/sym.rs
index 0931f604f..348f00915 100644
--- a/rust/src/sym.rs
+++ b/rust/src/sym.rs
@@ -26,6 +26,16 @@
// to the file of the Python module that defines it.
// tvm-ffi-stubgen(prefix): sym
// tvm-ffi-stubgen(skip): sym.ConstraintContext
+// The canonical simplifier's intermediate forms are internal to it.
+// tvm-ffi-stubgen(skip): sym.CanonicalExpr
+// tvm-ffi-stubgen(skip): sym.SplitExpr
+// tvm-ffi-stubgen(skip): sym.SumExpr
+// Integer sets and iterator sums are hand-written in their modules: C++
constructs them.
+// tvm-ffi-stubgen(skip): sym.IntSet
+// tvm-ffi-stubgen(skip): sym.IntervalSet
+// tvm-ffi-stubgen(skip): sym.IterMapExpr
+// tvm-ffi-stubgen(skip): sym.IterSplitExpr
+// tvm-ffi-stubgen(skip): sym.IterSumExpr
pub mod analyzer;
pub mod constraint;
diff --git a/rust/src/sym/analyzer.rs b/rust/src/sym/analyzer.rs
index 916cdb543..3bf325735 100644
--- a/rust/src/sym/analyzer.rs
+++ b/rust/src/sym/analyzer.rs
@@ -18,13 +18,11 @@
*/
//! The analyzer. An [`Analyzer`] proves, bounds and simplifies integer
expressions under the facts
-//! of a [`ConstraintContext`]; its bounds are [`ConstIntBound`]s and
[`ModularSet`]s, and
-//! [`SplitExpr`] and [`SumExpr`] are the forms its canonical simplifier works
in.
+//! of a [`ConstraintContext`]; its bounds are [`ConstIntBound`]s and
[`ModularSet`]s.
// tvm-ffi-stubgen(missing-docs): allow
-// Python has no class for `CanonicalExpr`, `SumExpr` and `SplitExpr`, the
forms of the canonical
-// simplifier, or for the `RewriteSimplifierStats` that
`Analyzer.rewrite_simplify_stats()`
-// returns; they are here, with the analyzer they belong to.
+// Python has no class for the `RewriteSimplifierStats` that
`Analyzer.rewrite_simplify_stats()`
+// returns; it is here, with the analyzer it belongs to.
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
@@ -40,17 +38,11 @@
// tvm-ffi-stubgen(import-object): crate::ir::support
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(import-object): crate::ir::expr::PrimExpr
-// The canonical forms are primitive expressions, as C++ `SplitExpr` and
`SumExpr` derive from
-// `PrimExpr`.
-// tvm-ffi-stubgen(upcast): sym.SplitExpr -> PrimExpr
-// tvm-ffi-stubgen(upcast): sym.SumExpr -> PrimExpr
// tvm-ffi-stubgen(import-object): std::borrow::Borrow
// tvm-ffi-stubgen(import-object): tvm_ffi::IntoObjectRef
// tvm-ffi-stubgen(begin): import-section
use crate::ExitOnDrop;
-use crate::ir::expr::Expr;
-use crate::ir::expr::ExprObj;
use crate::ir::expr::IntoExpr;
use crate::ir::expr::PrimExpr;
use crate::ir::expr::Var;
@@ -58,7 +50,6 @@ use crate::ir::support;
use crate::sym::constraint::ConstraintContext;
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;
@@ -583,102 +574,3 @@ impl RewriteSimplifierStats {
}
}
// tvm-ffi-stubgen(end)
-
-//
---------------------------------------------------------------------------------------------
-// The forms of the canonical simplifier
-
-/// The base of the intermediate forms the canonical simplifier introduces,
[`SplitExpr`] and
-/// [`SumExpr`].
-// tvm-ffi-stubgen(begin): object/sym.CanonicalExpr
-#[allow(missing_docs)]
-#[repr(C)]
-#[derive(tvm_ffi::derive::ObjectRef, Clone)]
-pub struct CanonicalExpr {
- data_: ObjectArc<CanonicalExprObj>,
-}
-
-/// The object behind [`CanonicalExpr`].
-///
-/// Opaque: no metadata of its own: total_size is unknown. Fields are read
through the C ABI getters.
-#[allow(missing_docs)]
-#[repr(C)]
-#[derive(tvm_ffi::derive::Object)]
-#[type_key = "sym.CanonicalExpr"]
-#[type_immutable]
-pub struct CanonicalExprObj {
- base_: ExprObj,
-}
-
-impl Deref for CanonicalExprObj {
- type Target = ExprObj;
- fn deref(&self) -> &ExprObj {
- &self.base_
- }
-}
-
-tvm_ffi::impl_object_ref_upcast!(CanonicalExpr => Expr);
-// tvm-ffi-stubgen(end)
-
-/// The canonical form of a split index: `((index % upper_factor) /
lower_factor) * scale`.
-// tvm-ffi-stubgen(begin): object/sym.SplitExpr
-#[allow(missing_docs)]
-#[repr(C)]
-#[derive(tvm_ffi::derive::ObjectRef, Clone)]
-pub struct SplitExpr {
- data_: ObjectArc<SplitExprObj>,
-}
-
-/// The object behind [`SplitExpr`].
-///
-/// Opaque: no metadata of its own: total_size is unknown. Fields are read
through the C ABI getters.
-#[allow(missing_docs)]
-#[repr(C)]
-#[derive(tvm_ffi::derive::Object)]
-#[type_key = "sym.SplitExpr"]
-#[type_immutable]
-pub struct SplitExprObj {
- base_: CanonicalExprObj,
-}
-
-impl Deref for SplitExprObj {
- type Target = CanonicalExprObj;
- fn deref(&self) -> &CanonicalExprObj {
- &self.base_
- }
-}
-
-tvm_ffi::impl_object_ref_upcast!(SplitExpr => Expr, SplitExpr =>
CanonicalExpr);
-tvm_ffi::impl_object_ref_upcast!(unsafe SplitExpr => PrimExpr);
-// tvm-ffi-stubgen(end)
-
-/// The canonical form of a sum: a sum of [`SplitExpr`]s plus a constant base.
-// tvm-ffi-stubgen(begin): object/sym.SumExpr
-#[allow(missing_docs)]
-#[repr(C)]
-#[derive(tvm_ffi::derive::ObjectRef, Clone)]
-pub struct SumExpr {
- data_: ObjectArc<SumExprObj>,
-}
-
-/// The object behind [`SumExpr`].
-///
-/// Opaque: no metadata of its own: total_size is unknown. Fields are read
through the C ABI getters.
-#[allow(missing_docs)]
-#[repr(C)]
-#[derive(tvm_ffi::derive::Object)]
-#[type_key = "sym.SumExpr"]
-#[type_immutable]
-pub struct SumExprObj {
- base_: CanonicalExprObj,
-}
-
-impl Deref for SumExprObj {
- type Target = CanonicalExprObj;
- fn deref(&self) -> &CanonicalExprObj {
- &self.base_
- }
-}
-
-tvm_ffi::impl_object_ref_upcast!(SumExpr => Expr, SumExpr => CanonicalExpr);
-tvm_ffi::impl_object_ref_upcast!(unsafe SumExpr => PrimExpr);
-// tvm-ffi-stubgen(end)
diff --git a/rust/src/sym/int_set.rs b/rust/src/sym/int_set.rs
index c9cc3ac8a..5182c7676 100644
--- a/rust/src/sym/int_set.rs
+++ b/rust/src/sym/int_set.rs
@@ -22,10 +22,17 @@
//! `estimate_region_*` functions bound the region a set of accesses covers.
// tvm-ffi-stubgen(missing-docs): allow
+// The sets are hand-written: C++ constructs them, so they have no field-wise
allocator.
+// tvm-ffi-stubgen(ty-map): sym.IntSet -> crate::sym::int_set::IntSet
+// tvm-ffi-stubgen(ty-map): sym.IntervalSet -> crate::sym::int_set::IntervalSet
+// tvm-ffi-stubgen(import-object): std::ops::Deref
+// tvm-ffi-stubgen(import-object): tvm_ffi::Object
+// tvm-ffi-stubgen(import-object): tvm_ffi::ObjectArc
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
// tvm-ffi-stubgen(import-object): tvm_ffi::Array
// tvm-ffi-stubgen(import-object): tvm_ffi::Map
+// tvm-ffi-stubgen(import-object): tvm_ffi::Result
// tvm-ffi-stubgen(import-object): crate::ir::expr::Range
// tvm-ffi-stubgen(import-object): crate::ir::expr::Var
// tvm-ffi-stubgen(import-object): crate::sym::analyzer::Analyzer
@@ -34,9 +41,6 @@
// tvm-ffi-stubgen(import-object): crate::ir::expr::PrimExpr
// tvm-ffi-stubgen(import-object): std::borrow::Borrow
// tvm-ffi-stubgen(import-object): tvm_ffi::IntoObjectRef
-// An interval is bounded by `PrimExpr`s, as C++ `IntervalSetObj` holds them.
-// tvm-ffi-stubgen(field): sym.IntervalSet.min_value -> PrimExpr
-// tvm-ffi-stubgen(field): sym.IntervalSet.max_value -> PrimExpr
// tvm-ffi-stubgen(begin): import-section
use crate::ir::expr::IntoExpr;
use crate::ir::expr::PrimExpr;
@@ -51,9 +55,7 @@ use tvm_ffi::IntoObjectRef;
use tvm_ffi::Map;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
-use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
-use tvm_ffi::reflection::FieldGetter;
// tvm-ffi-stubgen(end)
/// The base of every set of integers, one dimension each, that bounds the
values an expression
@@ -78,7 +80,6 @@ use tvm_ffi::reflection::FieldGetter;
/// assert_eq!(script(&bound)?, "sym.IntervalSet(min_value=neg_inf,
max_value=15)");
/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
/// ```
-// tvm-ffi-stubgen(begin): object/sym.IntSet
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::ObjectRef, Clone)]
@@ -88,15 +89,21 @@ pub struct IntSet {
/// The object behind [`IntSet`].
///
-/// Opaque: no metadata of its own: total_size is unknown. Fields are read
through the C ABI getters.
+/// Complete: reflected fields fill [24, 24) exactly.
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::Object)]
#[type_key = "sym.IntSet"]
+#[type_immutable]
+#[type_rs_parent]
pub struct IntSetObj {
base_: Object,
}
-// tvm-ffi-stubgen(end)
+
+const _: () = {
+ assert!(::core::mem::size_of::<IntSetObj>() == 24);
+ assert!(::core::mem::align_of::<IntSetObj>() == 8);
+};
impl IntSet {
/// Return whether the set is empty (`sym.IntSetIsNothing`).
@@ -127,7 +134,6 @@ impl IntSet {
/// The interval `[min_value, max_value]` of integers, whose bounds are
symbolic expressions and
/// may be [`neg_inf`] or [`pos_inf`].
-// tvm-ffi-stubgen(begin): object/sym.IntervalSet
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::ObjectRef, Clone)]
@@ -137,7 +143,7 @@ pub struct IntervalSet {
/// The object behind [`IntervalSet`].
///
-/// Opaque: parent 'sym.IntSet' is opaque (layout-unknown). Fields are read
through the C ABI getters.
+/// Complete: reflected fields fill [24, 40) exactly.
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::Object)]
@@ -146,8 +152,15 @@ pub struct IntervalSet {
#[type_immutable]
pub struct IntervalSetObj {
base_: IntSetObj,
+ pub min_value: PrimExpr,
+ pub max_value: PrimExpr,
}
+const _: () = {
+ assert!(::core::mem::size_of::<IntervalSetObj>() == 40);
+ assert!(::core::mem::align_of::<IntervalSetObj>() == 8);
+};
+
impl Deref for IntervalSetObj {
type Target = IntSetObj;
fn deref(&self) -> &IntSetObj {
@@ -155,19 +168,7 @@ impl Deref for IntervalSetObj {
}
}
-#[allow(missing_docs)]
-impl IntervalSetObj {
- pub fn min_value(&self) -> Result<PrimExpr> {
- FieldGetter::new(Self::type_index(), "min_value")?.get(self)
- }
-
- pub fn max_value(&self) -> Result<PrimExpr> {
- FieldGetter::new(Self::type_index(), "max_value")?.get(self)
- }
-}
-
tvm_ffi::impl_object_ref_upcast!(IntervalSet => IntSet);
-// tvm-ffi-stubgen(end)
impl IntervalSet {
/// Create the interval `[min_value, max_value]`, whose bounds may be
[`neg_inf`] and
diff --git a/rust/src/sym/iter_affine_map.rs b/rust/src/sym/iter_affine_map.rs
index ba0e00c6c..ab88dd731 100644
--- a/rust/src/sym/iter_affine_map.rs
+++ b/rust/src/sym/iter_affine_map.rs
@@ -22,10 +22,19 @@
//! functions simplify, divide and invert such maps.
// tvm-ffi-stubgen(missing-docs): allow
+// The iterator forms are hand-written: C++ constructs them, so they have no
field-wise
+// allocator.
+// tvm-ffi-stubgen(ty-map): sym.IterMapExpr ->
crate::sym::iter_affine_map::IterMapExpr
+// tvm-ffi-stubgen(ty-map): sym.IterSplitExpr ->
crate::sym::iter_affine_map::IterSplitExpr
+// tvm-ffi-stubgen(ty-map): sym.IterSumExpr ->
crate::sym::iter_affine_map::IterSumExpr
+// tvm-ffi-stubgen(import-object): crate::ir::expr::Expr
+// tvm-ffi-stubgen(import-object): crate::ir::expr::ExprObj
+// tvm-ffi-stubgen(import-object): std::ops::Deref
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
// tvm-ffi-stubgen(import-object): tvm_ffi::AnyView
// tvm-ffi-stubgen(import-object): tvm_ffi::Map
+// tvm-ffi-stubgen(import-object): tvm_ffi::Result
// tvm-ffi-stubgen(import-object): crate::ir::expr::Range
// tvm-ffi-stubgen(import-object): crate::ir::expr::Var
// tvm-ffi-stubgen(import-object): crate::sym::analyzer::Analyzer
@@ -33,21 +42,13 @@
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(import-object): crate::ir::expr::PrimExpr
// tvm-ffi-stubgen(import-object): crate::ir::expr::PrimVar
-// The iterator forms are primitive expressions, as C++ `IterMapExpr` derives
from `PrimExpr`, and
-// their marks and padding predicate are `PrimExpr`s.
-// tvm-ffi-stubgen(upcast): sym.IterMapExpr -> PrimExpr
-// tvm-ffi-stubgen(upcast): sym.IterSplitExpr -> PrimExpr
-// tvm-ffi-stubgen(upcast): sym.IterSumExpr -> PrimExpr
+// The marks and the padding predicate are `PrimExpr`s, as in C++.
// tvm-ffi-stubgen(field): sym.IterMark.source -> PrimExpr
// tvm-ffi-stubgen(field): sym.IterMark.extent -> PrimExpr
// tvm-ffi-stubgen(field): sym.IterMapResult.padding_predicate -> PrimExpr
// tvm-ffi-stubgen(nullable): sym.IterMapResult.padding_predicate
// tvm-ffi-stubgen(import-object): std::borrow::Borrow
// tvm-ffi-stubgen(import-object): tvm_ffi::IntoObjectRef
-// tvm-ffi-stubgen(field): sym.IterSplitExpr.lower_factor -> PrimExpr
-// tvm-ffi-stubgen(field): sym.IterSplitExpr.extent -> PrimExpr
-// tvm-ffi-stubgen(field): sym.IterSplitExpr.scale -> PrimExpr
-// tvm-ffi-stubgen(field): sym.IterSumExpr.base -> PrimExpr
// tvm-ffi-stubgen(begin): import-section
use crate::ir::expr::Expr;
use crate::ir::expr::ExprObj;
@@ -66,15 +67,12 @@ use tvm_ffi::IntoObjectRef;
use tvm_ffi::Map;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
-use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
use tvm_ffi::String;
-use tvm_ffi::reflection::FieldGetter;
// tvm-ffi-stubgen(end)
/// The base of the iterator map forms, [`IterSplitExpr`] and [`IterSumExpr`].
They hold the result
/// of detection and do not appear in a well-formed function.
-// tvm-ffi-stubgen(begin): object/sym.IterMapExpr
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::ObjectRef, Clone)]
@@ -84,16 +82,22 @@ pub struct IterMapExpr {
/// The object behind [`IterMapExpr`].
///
-/// Opaque: no metadata of its own: total_size is unknown. Fields are read
through the C ABI getters.
+/// Complete: reflected fields fill [40, 40) exactly.
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::Object)]
#[type_key = "sym.IterMapExpr"]
#[type_immutable]
+#[type_rs_parent]
pub struct IterMapExprObj {
base_: ExprObj,
}
+const _: () = {
+ assert!(::core::mem::size_of::<IterMapExprObj>() == 40);
+ assert!(::core::mem::align_of::<IterMapExprObj>() == 8);
+};
+
impl Deref for IterMapExprObj {
type Target = ExprObj;
fn deref(&self) -> &ExprObj {
@@ -103,7 +107,6 @@ impl Deref for IterMapExprObj {
tvm_ffi::impl_object_ref_upcast!(IterMapExpr => Expr);
tvm_ffi::impl_object_ref_upcast!(unsafe IterMapExpr => PrimExpr);
-// tvm-ffi-stubgen(end)
/// The result of [`detect_iter_map`]: the detected `indices`, or empty
`indices` and the `errors`
/// that say why detection failed.
@@ -237,7 +240,6 @@ impl IterMark {
// tvm-ffi-stubgen(end)
/// The split `(source // lower_factor) % extent * scale` of the marked
iterator `source`.
-// tvm-ffi-stubgen(begin): object/sym.IterSplitExpr
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::ObjectRef, Clone)]
@@ -247,7 +249,7 @@ pub struct IterSplitExpr {
/// The object behind [`IterSplitExpr`].
///
-/// Opaque: parent 'sym.IterMapExpr' is opaque (layout-unknown). Fields are
read through the C ABI getters.
+/// Complete: reflected fields fill [40, 72) exactly.
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::Object)]
@@ -256,8 +258,17 @@ pub struct IterSplitExpr {
#[type_immutable]
pub struct IterSplitExprObj {
base_: IterMapExprObj,
+ pub source: IterMark,
+ pub lower_factor: PrimExpr,
+ pub extent: PrimExpr,
+ pub scale: PrimExpr,
}
+const _: () = {
+ assert!(::core::mem::size_of::<IterSplitExprObj>() == 72);
+ assert!(::core::mem::align_of::<IterSplitExprObj>() == 8);
+};
+
impl Deref for IterSplitExprObj {
type Target = IterMapExprObj;
fn deref(&self) -> &IterMapExprObj {
@@ -265,31 +276,10 @@ impl Deref for IterSplitExprObj {
}
}
-#[allow(missing_docs)]
-impl IterSplitExprObj {
- pub fn source(&self) -> Result<IterMark> {
- FieldGetter::new(Self::type_index(), "source")?.get(self)
- }
-
- pub fn lower_factor(&self) -> Result<PrimExpr> {
- FieldGetter::new(Self::type_index(), "lower_factor")?.get(self)
- }
-
- pub fn extent(&self) -> Result<PrimExpr> {
- FieldGetter::new(Self::type_index(), "extent")?.get(self)
- }
-
- pub fn scale(&self) -> Result<PrimExpr> {
- FieldGetter::new(Self::type_index(), "scale")?.get(self)
- }
-}
-
tvm_ffi::impl_object_ref_upcast!(IterSplitExpr => Expr, IterSplitExpr =>
IterMapExpr);
tvm_ffi::impl_object_ref_upcast!(unsafe IterSplitExpr => PrimExpr);
-// tvm-ffi-stubgen(end)
/// The sum of the splits `args` plus `base`.
-// tvm-ffi-stubgen(begin): object/sym.IterSumExpr
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::ObjectRef, Clone)]
@@ -299,7 +289,7 @@ pub struct IterSumExpr {
/// The object behind [`IterSumExpr`].
///
-/// Opaque: parent 'sym.IterMapExpr' is opaque (layout-unknown). Fields are
read through the C ABI getters.
+/// Complete: reflected fields fill [40, 56) exactly.
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::Object)]
@@ -308,8 +298,15 @@ pub struct IterSumExpr {
#[type_immutable]
pub struct IterSumExprObj {
base_: IterMapExprObj,
+ pub args: Array<IterSplitExpr>,
+ pub base: PrimExpr,
}
+const _: () = {
+ assert!(::core::mem::size_of::<IterSumExprObj>() == 56);
+ assert!(::core::mem::align_of::<IterSumExprObj>() == 8);
+};
+
impl Deref for IterSumExprObj {
type Target = IterMapExprObj;
fn deref(&self) -> &IterMapExprObj {
@@ -317,20 +314,8 @@ impl Deref for IterSumExprObj {
}
}
-#[allow(missing_docs)]
-impl IterSumExprObj {
- pub fn args(&self) -> Result<Array<IterSplitExpr>> {
- FieldGetter::new(Self::type_index(), "args")?.get(self)
- }
-
- pub fn base(&self) -> Result<PrimExpr> {
- FieldGetter::new(Self::type_index(), "base")?.get(self)
- }
-}
-
tvm_ffi::impl_object_ref_upcast!(IterSumExpr => Expr, IterSumExpr =>
IterMapExpr);
tvm_ffi::impl_object_ref_upcast!(unsafe IterSumExpr => PrimExpr);
-// tvm-ffi-stubgen(end)
// `IterMark::new` is generated: the native `sym.IterMark` constructor stores
the fields as given.
diff --git a/rust/src/target.rs b/rust/src/target.rs
index 6085fa3f4..a62fc2e11 100644
--- a/rust/src/target.rs
+++ b/rust/src/target.rs
@@ -24,12 +24,19 @@
// tvm-ffi-stubgen(missing-docs): allow
// Every object registered under `target` has its block in this file.
// tvm-ffi-stubgen(prefix): target
+// `target.Target` is hand-written below: a target is built by parsing its
config, never from
+// raw fields.
+// tvm-ffi-stubgen(skip): target.Target
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
// tvm-ffi-stubgen(import-object): tvm_ffi::Module
// tvm-ffi-stubgen(import-object): crate::ir::module::IRModule
// tvm-ffi-stubgen(import-object): tvm_ffi::IntoObjectRef
// tvm-ffi-stubgen(import-object): std::borrow::Borrow
+// tvm-ffi-stubgen(ty-map): target.Target -> crate::target::Target
+// tvm-ffi-stubgen(import-object): tvm_ffi::Any
+// tvm-ffi-stubgen(import-object): tvm_ffi::Map
+// tvm-ffi-stubgen(import-object): tvm_ffi::object::ObjectRef
// tvm-ffi-stubgen(begin): import-section
use crate::ir::module::IRModule;
@@ -44,6 +51,7 @@ use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
use tvm_ffi::String;
+use tvm_ffi::extra::cold_init_lazy_lock::ColdInitLazyLock;
use tvm_ffi::object::ObjectRef;
use tvm_ffi::reflection::FieldGetter;
// tvm-ffi-stubgen(end)
@@ -67,7 +75,6 @@ use tvm_ffi::reflection::FieldGetter;
/// assert!(Target::current(true)?.is_none());
/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
/// ```
-// tvm-ffi-stubgen(begin): object/target.Target
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::ObjectRef, Clone)]
@@ -77,7 +84,7 @@ pub struct Target {
/// The object behind [`Target`].
///
-/// Opaque: bytes [72, 88) of [24, 88) are not accounted for by reflected
fields. Fields are read through the C ABI getters.
+/// Complete: reflected fields fill [24, 72) exactly.
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::Object)]
@@ -85,31 +92,17 @@ pub struct Target {
#[type_final]
pub struct TargetObj {
base_: Object,
+ pub kind: TargetKind,
+ pub host: Option<ObjectRef>,
+ pub tag: String,
+ pub keys: Array<String>,
+ pub attrs: Map<String, Any>,
}
-#[allow(missing_docs)]
-impl TargetObj {
- pub fn kind(&self) -> Result<TargetKind> {
- FieldGetter::new(Self::type_index(), "kind")?.get(self)
- }
-
- pub fn tag(&self) -> Result<String> {
- FieldGetter::new(Self::type_index(), "tag")?.get(self)
- }
-
- pub fn keys(&self) -> Result<Array<String>> {
- FieldGetter::new(Self::type_index(), "keys")?.get(self)
- }
-
- pub fn attrs(&self) -> Result<Map<String, Any>> {
- FieldGetter::new(Self::type_index(), "attrs")?.get(self)
- }
-
- pub fn host(&self) -> Result<Option<ObjectRef>> {
- FieldGetter::new(Self::type_index(), "host")?.get(self)
- }
-}
-// tvm-ffi-stubgen(end)
+const _: () = {
+ assert!(::core::mem::size_of::<TargetObj>() == 72);
+ assert!(::core::mem::align_of::<TargetObj>() == 8);
+};
impl Target {
/// Create the target that a tag (`"nvidia/nvidia-a100"`), a kind name
(`"llvm"`) or a JSON
@@ -232,15 +225,22 @@ pub struct TargetKindObj {
#[allow(missing_docs)]
impl TargetKindObj {
pub fn name(&self) -> Result<String> {
- FieldGetter::new(Self::type_index(), "name")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(TargetKindObj::type_index(), "name"));
+ GETTER.clone()?.get(self)
}
pub fn default_device_type(&self) -> Result<i64> {
- FieldGetter::new(Self::type_index(), "default_device_type")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
ColdInitLazyLock::new(|| {
+ FieldGetter::new(TargetKindObj::type_index(),
"default_device_type")
+ });
+ GETTER.clone()?.get(self)
}
pub fn default_keys(&self) -> Result<Array<String>> {
- FieldGetter::new(Self::type_index(), "default_keys")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(TargetKindObj::type_index(), "default_keys"));
+ GETTER.clone()?.get(self)
}
}
// tvm-ffi-stubgen(end)
@@ -285,7 +285,7 @@ pub fn list_tags() -> Result<Map<String, Target>> {
/// let config: Map<String, Any> =
/// [("kind".into(),
Any::from(String::from("llvm")))].into_iter().collect();
/// target::register_tag("my_ext/cpu", config, false)?;
-/// assert_eq!(Target::new("my_ext/cpu")?.kind()?.name()?, "llvm");
+/// assert_eq!(Target::new("my_ext/cpu")?.kind.name()?, "llvm");
/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
/// ```
pub fn register_tag(
diff --git a/rust/src/tirx.rs b/rust/src/tirx.rs
index d0607d711..1b8918897 100644
--- a/rust/src/tirx.rs
+++ b/rust/src/tirx.rs
@@ -25,6 +25,10 @@
// An object registered under `tirx` that no file has a block for yet is
generated here, to be moved
// to the file of the Python module that defines it.
// tvm-ffi-stubgen(prefix): tirx
+// Layouts are hand-written in `layout`: C++ constructs them.
+// tvm-ffi-stubgen(skip): tirx.Layout
+// tvm-ffi-stubgen(skip): tirx.TileLayout
+// tvm-ffi-stubgen(skip): tirx.ComposeLayout
pub mod analysis;
pub mod function;
diff --git a/rust/src/tirx/layout.rs b/rust/src/tirx/layout.rs
index 4b9f61a64..f70ae26e7 100644
--- a/rust/src/tirx/layout.rs
+++ b/rust/src/tirx/layout.rs
@@ -22,12 +22,19 @@
//! strided [`Iter`] factors, and a [`ComposeLayout`] swizzles the memory
coordinate of one.
// tvm-ffi-stubgen(missing-docs): allow
+// The layouts are hand-written: C++ constructs them, so they have no
field-wise allocator.
+// tvm-ffi-stubgen(ty-map): tirx.Layout -> crate::tirx::layout::Layout
+// tvm-ffi-stubgen(ty-map): tirx.TileLayout -> crate::tirx::layout::TileLayout
+// tvm-ffi-stubgen(ty-map): tirx.ComposeLayout ->
crate::tirx::layout::ComposeLayout
+// tvm-ffi-stubgen(import-object): std::ops::Deref
+// tvm-ffi-stubgen(import-object): tvm_ffi::Array
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
// tvm-ffi-stubgen(import-object): crate::sym::analyzer::Analyzer
// tvm-ffi-stubgen(import-object): crate::tirx::thread_scope::ThreadScope
// tvm-ffi-stubgen(import-object): tvm_ffi::Any
// tvm-ffi-stubgen(import-object): tvm_ffi::String
+// tvm-ffi-stubgen(import-object): tvm_ffi::Result
// tvm-ffi-stubgen(import-object): crate::ir::support
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(import-object): std::borrow::Borrow
@@ -35,10 +42,9 @@
// An iterator's extent and stride are `PrimExpr`s, as C++ `IterObj` holds
them.
// tvm-ffi-stubgen(field): tirx.Iter.extent -> PrimExpr
// tvm-ffi-stubgen(field): tirx.Iter.stride -> PrimExpr
-// A tile layout's offsets are `PrimExpr`s, as C++ `TileLayoutObj` holds them.
-// tvm-ffi-stubgen(field): tirx.TileLayout.offset -> Map<Axis, PrimExpr>
// tvm-ffi-stubgen(import-object): tvm_ffi::Map
// tvm-ffi-stubgen(import-object): tvm_ffi::IntoObjectRef
+
// tvm-ffi-stubgen(begin): import-section
use crate::ir::expr::IntoExpr;
use crate::ir::expr::PrimExpr;
@@ -53,11 +59,9 @@ use tvm_ffi::IntoObjectRef;
use tvm_ffi::Map;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
-use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
use tvm_ffi::String;
use tvm_ffi::enum_ref::EnumObj;
-use tvm_ffi::reflection::FieldGetter;
// tvm-ffi-stubgen(end)
/// The base of every layout: a map from the logical indices of a tensor to a
coordinate on each
@@ -87,7 +91,6 @@ use tvm_ffi::reflection::FieldGetter;
/// assert_eq!(script(&offset)?, "i * 8 + j");
/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
/// ```
-// tvm-ffi-stubgen(begin): object/tirx.Layout
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::ObjectRef, Clone)]
@@ -97,15 +100,21 @@ pub struct Layout {
/// The object behind [`Layout`].
///
-/// Opaque: no metadata of its own: total_size is unknown. Fields are read
through the C ABI getters.
+/// Complete: reflected fields fill [24, 24) exactly.
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::Object)]
#[type_key = "tirx.Layout"]
+#[type_immutable]
+#[type_rs_parent]
pub struct LayoutObj {
base_: Object,
}
-// tvm-ffi-stubgen(end)
+
+const _: () = {
+ assert!(::core::mem::size_of::<LayoutObj>() == 24);
+ assert!(::core::mem::align_of::<LayoutObj>() == 8);
+};
impl Layout {
/// Check the invariants of the layout (`tirx.LayoutCheckWellFormed`).
@@ -222,7 +231,6 @@ impl Layout {
/// The layout of strided factors: each [`Iter`] of `shard` spans distinct
elements and each of
/// `replica` their copies, adding its stride per position to the coordinate
on its axis, from the
/// origin `offset` gives on each axis.
-// tvm-ffi-stubgen(begin): object/tirx.TileLayout
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::ObjectRef, Clone)]
@@ -232,7 +240,7 @@ pub struct TileLayout {
/// The object behind [`TileLayout`].
///
-/// Opaque: parent 'tirx.Layout' is opaque (layout-unknown). Fields are read
through the C ABI getters.
+/// Complete: reflected fields fill [24, 48) exactly.
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::Object)]
@@ -241,8 +249,16 @@ pub struct TileLayout {
#[type_immutable]
pub struct TileLayoutObj {
base_: LayoutObj,
+ pub shard: Array<Iter>,
+ pub replica: Array<Iter>,
+ pub offset: Map<Axis, PrimExpr>,
}
+const _: () = {
+ assert!(::core::mem::size_of::<TileLayoutObj>() == 48);
+ assert!(::core::mem::align_of::<TileLayoutObj>() == 8);
+};
+
impl Deref for TileLayoutObj {
type Target = LayoutObj;
fn deref(&self) -> &LayoutObj {
@@ -250,23 +266,7 @@ impl Deref for TileLayoutObj {
}
}
-#[allow(missing_docs)]
-impl TileLayoutObj {
- pub fn shard(&self) -> Result<Array<Iter>> {
- FieldGetter::new(Self::type_index(), "shard")?.get(self)
- }
-
- pub fn replica(&self) -> Result<Array<Iter>> {
- FieldGetter::new(Self::type_index(), "replica")?.get(self)
- }
-
- pub fn offset(&self) -> Result<Map<Axis, PrimExpr>> {
- FieldGetter::new(Self::type_index(), "offset")?.get(self)
- }
-}
-
tvm_ffi::impl_object_ref_upcast!(TileLayout => Layout);
-// tvm-ffi-stubgen(end)
impl TileLayout {
/// Create the tile layout of the `shard` and `replica` iterators, shifted
by `offset`
@@ -467,7 +467,6 @@ impl Axis {
/// The layout `tile_layout` with its memory coordinate XOR-swizzled: above
the low `per_element`
/// bits, the group of `swizzle_len` bits `atom_len` bits up is XORed into the
group below if
/// `swizzle_inner`, and the lower group into the upper one otherwise.
-// tvm-ffi-stubgen(begin): object/tirx.ComposeLayout
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::ObjectRef, Clone)]
@@ -477,7 +476,7 @@ pub struct ComposeLayout {
/// The object behind [`ComposeLayout`].
///
-/// Opaque: parent 'tirx.Layout' is opaque (layout-unknown). Fields are read
through the C ABI getters.
+/// Complete: reflected fields fill [24, 56) exactly (alignment padding [37,
40)).
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::Object)]
@@ -486,8 +485,20 @@ pub struct ComposeLayout {
#[type_immutable]
pub struct ComposeLayoutObj {
base_: LayoutObj,
+ pub per_element: i32,
+ pub swizzle_len: i32,
+ pub atom_len: i32,
+ pub swizzle_inner: bool,
+ pub tile_layout: TileLayout,
+ pub inner_mask: i32,
+ pub outer_mask: i32,
}
+const _: () = {
+ assert!(::core::mem::size_of::<ComposeLayoutObj>() == 56);
+ assert!(::core::mem::align_of::<ComposeLayoutObj>() == 8);
+};
+
impl Deref for ComposeLayoutObj {
type Target = LayoutObj;
fn deref(&self) -> &LayoutObj {
@@ -495,39 +506,7 @@ impl Deref for ComposeLayoutObj {
}
}
-#[allow(missing_docs)]
-impl ComposeLayoutObj {
- pub fn per_element(&self) -> Result<i64> {
- FieldGetter::new(Self::type_index(), "per_element")?.get(self)
- }
-
- pub fn swizzle_len(&self) -> Result<i64> {
- FieldGetter::new(Self::type_index(), "swizzle_len")?.get(self)
- }
-
- pub fn atom_len(&self) -> Result<i64> {
- FieldGetter::new(Self::type_index(), "atom_len")?.get(self)
- }
-
- pub fn swizzle_inner(&self) -> Result<bool> {
- FieldGetter::new(Self::type_index(), "swizzle_inner")?.get(self)
- }
-
- pub fn inner_mask(&self) -> Result<i64> {
- FieldGetter::new(Self::type_index(), "inner_mask")?.get(self)
- }
-
- pub fn outer_mask(&self) -> Result<i64> {
- FieldGetter::new(Self::type_index(), "outer_mask")?.get(self)
- }
-
- pub fn tile_layout(&self) -> Result<TileLayout> {
- FieldGetter::new(Self::type_index(), "tile_layout")?.get(self)
- }
-}
-
tvm_ffi::impl_object_ref_upcast!(ComposeLayout => Layout);
-// tvm-ffi-stubgen(end)
impl ComposeLayout {
/// Create `tile_layout` with its memory coordinate swizzled: above the
low `per_element` bits,
diff --git a/rust/src/tirx/structural.rs b/rust/src/tirx/structural.rs
index 4aeeb8a57..136e5ff0c 100644
--- a/rust/src/tirx/structural.rs
+++ b/rust/src/tirx/structural.rs
@@ -29,6 +29,7 @@ use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
+use tvm_ffi::extra::cold_init_lazy_lock::ColdInitLazyLock;
use tvm_ffi::reflection::FieldGetter;
/// Header-only stand-in for the builtin `ffi.StructuralMutator`; it only
carries the ancestor depth.
@@ -77,7 +78,10 @@ pub struct StructuralMutatorWithConstraintContextObj {
#[allow(missing_docs)]
impl StructuralMutatorWithConstraintContextObj {
pub fn constraints(&self) -> Result<ConstraintContext> {
- FieldGetter::new(Self::type_index(), "constraints")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
ColdInitLazyLock::new(|| {
+
FieldGetter::new(StructuralMutatorWithConstraintContextObj::type_index(),
"constraints")
+ });
+ GETTER.clone()?.get(self)
}
}
// tvm-ffi-stubgen(end)
@@ -107,7 +111,10 @@ pub struct StructuralVisitorWithConstraintContextObj {
#[allow(missing_docs)]
impl StructuralVisitorWithConstraintContextObj {
pub fn constraints(&self) -> Result<ConstraintContext> {
- FieldGetter::new(Self::type_index(), "constraints")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
ColdInitLazyLock::new(|| {
+
FieldGetter::new(StructuralVisitorWithConstraintContextObj::type_index(),
"constraints")
+ });
+ GETTER.clone()?.get(self)
}
}
// tvm-ffi-stubgen(end)
diff --git a/rust/src/tirx/tile/dispatch.rs b/rust/src/tirx/tile/dispatch.rs
index af079ff18..d8deb5b64 100644
--- a/rust/src/tirx/tile/dispatch.rs
+++ b/rust/src/tirx/tile/dispatch.rs
@@ -58,6 +58,7 @@ use tvm_ffi::ObjectCore;
use tvm_ffi::ObjectRefCore;
use tvm_ffi::Result;
use tvm_ffi::String;
+use tvm_ffi::extra::cold_init_lazy_lock::ColdInitLazyLock;
use tvm_ffi::object::ObjectRef;
use tvm_ffi::reflection::FieldGetter;
// tvm-ffi-stubgen(end)
@@ -118,43 +119,70 @@ pub struct DispatchContextObj {
#[allow(missing_docs)]
impl DispatchContextObj {
pub fn target(&self) -> Result<Target> {
- FieldGetter::new(Self::type_index(), "target")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(DispatchContextObj::type_index(), "target"));
+ GETTER.clone()?.get(self)
}
pub fn thread_scope(&self) -> Result<ThreadScope> {
- FieldGetter::new(Self::type_index(), "thread_scope")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
ColdInitLazyLock::new(|| {
+ FieldGetter::new(DispatchContextObj::type_index(), "thread_scope")
+ });
+ GETTER.clone()?.get(self)
}
pub fn launch_params(&self) -> Result<Any> {
- FieldGetter::new(Self::type_index(), "launch_params")?.get_any(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
ColdInitLazyLock::new(|| {
+ FieldGetter::new(DispatchContextObj::type_index(), "launch_params")
+ });
+ GETTER.clone()?.get_any(self)
}
pub fn var_range_map(&self) -> Result<Map<Var, Range>> {
- FieldGetter::new(Self::type_index(), "var_range_map")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
ColdInitLazyLock::new(|| {
+ FieldGetter::new(DispatchContextObj::type_index(), "var_range_map")
+ });
+ GETTER.clone()?.get(self)
}
pub fn alloc_only(&self) -> Result<bool> {
- FieldGetter::new(Self::type_index(), "alloc_only")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
ColdInitLazyLock::new(|| {
+ FieldGetter::new(DispatchContextObj::type_index(), "alloc_only")
+ });
+ GETTER.clone()?.get(self)
}
pub fn callbacks(&self) -> Result<Map<String, ObjectRef>> {
- FieldGetter::new(Self::type_index(), "callbacks")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
ColdInitLazyLock::new(|| {
+ FieldGetter::new(DispatchContextObj::type_index(), "callbacks")
+ });
+ GETTER.clone()?.get(self)
}
pub fn shared_state(&self) -> Result<Map<String, ObjectRef>> {
- FieldGetter::new(Self::type_index(), "shared_state")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
ColdInitLazyLock::new(|| {
+ FieldGetter::new(DispatchContextObj::type_index(), "shared_state")
+ });
+ GETTER.clone()?.get(self)
}
pub fn inter(&self) -> Result<Map<String, Array<PrimExpr>>> {
- FieldGetter::new(Self::type_index(), "inter")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(DispatchContextObj::type_index(), "inter"));
+ GETTER.clone()?.get(self)
}
pub fn intra(&self) -> Result<Map<String, Array<PrimExpr>>> {
- FieldGetter::new(Self::type_index(), "intra")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(DispatchContextObj::type_index(), "intra"));
+ GETTER.clone()?.get(self)
}
pub fn scope_kind(&self) -> Result<String> {
- FieldGetter::new(Self::type_index(), "scope_kind")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
ColdInitLazyLock::new(|| {
+ FieldGetter::new(DispatchContextObj::type_index(), "scope_kind")
+ });
+ GETTER.clone()?.get(self)
}
}
// tvm-ffi-stubgen(end)
diff --git a/rust/src/tirx/transform/function_pass.rs
b/rust/src/tirx/transform/function_pass.rs
index bef18cdec..3d0101c15 100644
--- a/rust/src/tirx/transform/function_pass.rs
+++ b/rust/src/tirx/transform/function_pass.rs
@@ -45,6 +45,7 @@ use tvm_ffi::IntoObjectRef;
use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
+use tvm_ffi::extra::cold_init_lazy_lock::ColdInitLazyLock;
use tvm_ffi::reflection::FieldGetter;
use tvm_ffi::rvalue_ref::RValueRef;
// tvm-ffi-stubgen(end)
@@ -107,7 +108,9 @@ impl Deref for FunctionPassObj {
#[allow(missing_docs)]
impl FunctionPassObj {
pub fn pass_info(&self) -> Result<PassInfo> {
- FieldGetter::new(Self::type_index(), "pass_info")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(FunctionPassObj::type_index(), "pass_info"));
+ GETTER.clone()?.get(self)
}
}
diff --git a/rust/tests/sym/analyzer.rs b/rust/tests/sym/analyzer.rs
index 0e68fb34b..ad15c6b41 100644
--- a/rust/tests/sym/analyzer.rs
+++ b/rust/tests/sym/analyzer.rs
@@ -19,7 +19,7 @@
//! Tests of `sym::analyzer`: the analyzer and its results.
-use tvm_tirx::ir::{Expr, ExprObj};
+use tvm_tirx::ir::Expr;
use tvm_tirx::sym::*;
use tvm_tirx::tvm_ffi::Object;
@@ -31,12 +31,8 @@ fn generated_types_match_the_registry() {
assert_type_contracts!(true;
ConstIntBoundObj: Object, ModularSetObj: Object,
RewriteSimplifierStatsObj: Object,
);
- // The analyzer holds its sub-analyzers unreflected; the expression
hierarchy registers no
- // layout of its own, so the canonical expressions are opaque too.
- assert_type_contracts!(false;
- AnalyzerObj: Object, CanonicalExprObj: ExprObj, SplitExprObj:
CanonicalExprObj,
- SumExprObj: CanonicalExprObj,
- );
+ // The analyzer holds its sub-analyzers unreflected.
+ assert_type_contracts!(false; AnalyzerObj: Object);
}
#[test]
diff --git a/rust/tests/sym/int_set.rs b/rust/tests/sym/int_set.rs
index 1796d4b47..b2e6065b3 100644
--- a/rust/tests/sym/int_set.rs
+++ b/rust/tests/sym/int_set.rs
@@ -49,15 +49,15 @@ fn interval_sets_are_read_through_accessors() {
.unwrap();
let interval: IntervalSet = set.downcast().unwrap();
let as_int = |expr: PrimExpr|
expr.downcast::<IntImm>().unwrap().value.to_i64().unwrap();
- assert_eq!(as_int(interval.min_value().unwrap()), 2);
- assert_eq!(as_int(interval.max_value().unwrap()), 20);
+ assert_eq!(as_int(interval.min_value.clone()), 2);
+ assert_eq!(as_int(interval.max_value.clone()), 20);
}
#[test]
fn int_sets() {
let analyzer = Analyzer::new();
let interval = IntervalSet::new(int32(0), int32(10));
- assert!(analyzer.can_prove_equal(interval.min_value().unwrap(),
int32(0)).unwrap());
+ assert!(analyzer.can_prove_equal(interval.min_value.clone(),
int32(0)).unwrap());
let set: IntSet = interval.into();
assert!(!set.is_nothing().unwrap());
assert!(!set.is_everything().unwrap());
diff --git a/rust/tests/sym/iter_affine_map.rs
b/rust/tests/sym/iter_affine_map.rs
index a89781950..7a48a4c62 100644
--- a/rust/tests/sym/iter_affine_map.rs
+++ b/rust/tests/sym/iter_affine_map.rs
@@ -67,22 +67,22 @@ fn iter_maps_decompose_into_marks_and_splits() {
// Each index is one split of a mark of extent 4, whose source is itself a
sum over the
// mark of x on [0, 16): x // 4 takes lower factor 4 of it, x % 4 lower
factor 1.
let only_split = |sum: &IterSumExpr| -> IterSplitExpr {
- let splits = sum.args().unwrap();
+ let splits = sum.args.clone();
assert_eq!(splits.len(), 1);
splits.get(0).unwrap()
};
let mut lower_factors = Vec::new();
for sum in result.indices.iter() {
- assert_eq!(as_int(sum.base().unwrap()), 0);
+ assert_eq!(as_int(sum.base.clone()), 0);
let outer = only_split(sum);
- let mark: IterMark = outer.source().unwrap();
+ let mark: IterMark = outer.source.clone();
assert_eq!(as_int(mark.extent.clone()), 4);
let inner_sum: IterSumExpr = mark.source.clone().downcast().unwrap();
let inner = only_split(&inner_sum);
- let x_mark: IterMark = inner.source().unwrap();
+ let x_mark: IterMark = inner.source.clone();
assert!(mark_is(&x_mark, &x, 16));
- assert_eq!(as_int(inner.extent().unwrap()), 4);
- lower_factors.push(as_int(inner.lower_factor().unwrap()));
+ assert_eq!(as_int(inner.extent.clone()), 4);
+ lower_factors.push(as_int(inner.lower_factor.clone()));
}
assert_eq!(lower_factors, vec![4, 1]);
}
@@ -105,7 +105,7 @@ fn iter_maps() {
assert!(analyzer.can_prove_equal(normalized, index.clone()).unwrap());
let sum = normalize_to_iter_sum(index.clone(), &dom,
Some(&analyzer)).unwrap();
- assert_eq!(sum.args().unwrap().len(), 2);
+ assert_eq!(sum.args.len(), 2);
let simplified =
iter_map_simplify([index.clone()], &dom, &predicate,
IterMapLevel::default(), None, true)
@@ -130,7 +130,7 @@ fn iter_maps() {
let mark = IterMark::new(x.clone(), int32(4));
let split = IterSplitExpr::new(mark, int32(1), int32(4),
int32(1)).unwrap();
- assert!(analyzer.can_prove_equal(split.scale().unwrap(),
int32(1)).unwrap());
+ assert!(analyzer.can_prove_equal(split.scale.clone(), int32(1)).unwrap());
let sum = IterSumExpr::new([split], int32(0));
- assert_eq!(sum.args().unwrap().len(), 1);
+ assert_eq!(sum.args.len(), 1);
}
diff --git a/rust/tests/target/main.rs b/rust/tests/target/main.rs
index c89b466b5..56bed5479 100644
--- a/rust/tests/target/main.rs
+++ b/rust/tests/target/main.rs
@@ -46,23 +46,22 @@ fn generated_types_match_the_registry() {
#[test]
fn native_targets_are_read_through_accessors() {
let llvm = native_target("llvm");
- let kind: TargetKind = llvm.kind().unwrap();
+ let kind: TargetKind = llvm.kind.clone();
assert_eq!(kind.name().unwrap().as_str(), "llvm");
assert!(kind.default_keys().unwrap().iter().any(|key| key.as_str() ==
"cpu"));
- assert!(llvm.keys().unwrap().iter().any(|key| key.as_str() == "cpu"));
- assert!(llvm.host().unwrap().is_none());
+ assert!(llvm.keys.iter().any(|key| key.as_str() == "cpu"));
+ assert!(llvm.host.is_none());
// The kind's reflected device type is what the native query reports.
let cuda = native_target("cuda");
- assert_eq!(cuda.kind().unwrap().name().unwrap().as_str(), "cuda");
+ assert_eq!(cuda.kind.name().unwrap().as_str(), "cuda");
let device_type: i64 = global("target.TargetGetDeviceType")
.call_tuple((cuda.clone(),))
.unwrap()
.try_into()
.unwrap();
- assert_eq!(cuda.kind().unwrap().default_device_type().unwrap(),
device_type);
+ assert_eq!(cuda.kind.default_device_type().unwrap(), device_type);
assert_ne!(device_type, kind.default_device_type().unwrap());
- cuda.attrs().unwrap();
// `host` is declared as an `ObjectRef`; it holds a `Target`.
let with_host: Target = global("target.WithHost")
@@ -70,28 +69,28 @@ fn native_targets_are_read_through_accessors() {
.unwrap()
.try_into()
.unwrap();
- let host: Target = with_host.host().unwrap().unwrap().downcast().unwrap();
- assert_eq!(host.kind().unwrap().name().unwrap().as_str(), "llvm");
-
assert!(with_host.host().unwrap().unwrap().downcast::<TargetKind>().is_err());
+ let host: Target = with_host.host.clone().unwrap().downcast().unwrap();
+ assert_eq!(host.kind.name().unwrap().as_str(), "llvm");
+ assert!(with_host.host.clone().unwrap().downcast::<TargetKind>().is_err());
}
#[test]
fn construction_and_queries() {
let llvm = Target::new("llvm").unwrap();
- assert_eq!(llvm.kind().unwrap().name().unwrap().as_str(), "llvm");
+ assert_eq!(llvm.kind.name().unwrap().as_str(), "llvm");
assert_eq!(llvm.get_target_device_type().unwrap(), 1);
assert!(llvm.export().unwrap().contains_key(&FfiString::from("kind")).unwrap());
assert!(llvm.as_json().unwrap().as_str().contains("llvm"));
assert!(llvm.feature("no_such_feature").unwrap().is_none());
let c = Target::from_config(config("c")).unwrap();
- assert_eq!(c.kind().unwrap().name().unwrap().as_str(), "c");
- assert!(llvm.with_host(Some(&c)).unwrap().host().unwrap().is_some());
- assert!(llvm.with_host(None).unwrap().host().unwrap().is_none());
+ assert_eq!(c.kind.name().unwrap().as_str(), "c");
+ assert!(llvm.with_host(Some(&c)).unwrap().host.is_some());
+ assert!(llvm.with_host(None).unwrap().host.is_none());
let kinds = Target::list_kinds().unwrap();
assert!(kinds.iter().any(|kind| kind.as_str() == "llvm"));
- let options = llvm.kind().unwrap().options().unwrap();
+ let options = llvm.kind.options().unwrap();
assert!(options.contains_key(&FfiString::from("mtriple")).unwrap());
assert_eq!(TargetKind::options_from_name("llvm").unwrap().len(),
options.len());
}
@@ -103,7 +102,7 @@ fn scope_ends_on_drop() {
{
let _scope = Target::new("c").unwrap().scope().unwrap();
let current = Target::current(false).unwrap().unwrap();
- assert_eq!(current.kind().unwrap().name().unwrap().as_str(), "c");
+ assert_eq!(current.kind.name().unwrap().as_str(), "c");
}
assert!(Target::current(true).unwrap().is_none());
}
@@ -111,7 +110,7 @@ fn scope_ends_on_drop() {
#[test]
fn tags() {
let tagged = target::register_tag("rust-test/llvm", config("llvm"),
true).unwrap();
- assert_eq!(tagged.kind().unwrap().name().unwrap().as_str(), "llvm");
+ assert_eq!(tagged.kind.name().unwrap().as_str(), "llvm");
let tags = target::list_tags().unwrap();
assert!(tags.contains_key(&FfiString::from("rust-test/llvm")).unwrap());
assert!(target::register_tag("rust-test/llvm", config("llvm"),
false).is_err());
diff --git a/rust/tests/tirx/layout.rs b/rust/tests/tirx/layout.rs
index a55ac3b80..20cd8b431 100644
--- a/rust/tests/tirx/layout.rs
+++ b/rust/tests/tirx/layout.rs
@@ -65,7 +65,7 @@ fn layouts() {
// Memory offset 10 becomes 10 ^ ((10 >> 3) & 3) = 11.
let swizzled = ComposeLayout::new(0, 2, 3, &tile, true).unwrap();
- assert_eq!(swizzled.atom_len().unwrap(), 3);
+ assert_eq!(swizzled.atom_len, 3);
assert_eq!(at(Layout::from(&swizzled).apply([int32(1), int32(2)],
&analyzer).unwrap()), 11);
}
diff --git a/src/backend/cuda/transform/bind_backend_config.cc
b/src/backend/cuda/transform/bind_backend_config.cc
index 13223a4ce..bdab846c9 100644
--- a/src/backend/cuda/transform/bind_backend_config.cc
+++ b/src/backend/cuda/transform/bind_backend_config.cc
@@ -20,13 +20,15 @@
/*! \file bind_backend_config.cc
* \brief Resolve CUDA configuration before any target-sensitive lowering.
*/
+#include <tvm/ffi/extra/structural_equal.h>
+#include <tvm/ffi/extra/structural_key.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/unique_name_supply.h>
#include <tvm/target/target.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>
-#include <map>
+#include <unordered_map>
#include <unordered_set>
#include <vector>
@@ -156,7 +158,7 @@ class DeviceHelperSpecializer : public StmtExprMutator {
if (auto existing = replacements_.Get(callee.value())) {
replacement = existing.value();
} else {
- bool same_group = original_target.value()->str() == target_->str() &&
+ bool same_group = ffi::StructuralEqual()(original_target.value(),
target_) &&
original->GetAttr<ffi::String>("backend_config").value_or("") == config_;
replacement = same_group
? callee.value()
@@ -239,7 +241,7 @@ Pass SpecializeDeviceHelpers() {
IRModule result({}, mod->attrs, mod->global_infos);
UniqueNameSupply names(mod->functions.begin(), mod->functions.end(),
[](const auto& kv) { return kv.first->name_hint; });
- std::map<std::string, ffi::ObjectPtr<DeviceHelperSpecializer>> groups;
+ std::unordered_map<ffi::StructuralKey,
ffi::ObjectPtr<DeviceHelperSpecializer>> groups;
for (auto [gv, base] : mod->functions) {
auto func = base.as<Function>();
if (!func.has_value()) {
@@ -257,8 +259,7 @@ Pass SpecializeDeviceHelpers() {
continue;
}
auto config =
func.value()->GetAttr<ffi::String>("backend_config").value_or("");
- std::string key = std::string(target.value()->str()) + "\n" +
std::string(config);
- auto& group = groups[key];
+ auto& group =
groups[ffi::StructuralKey(ffi::Array<ffi::Any>{target.value(), config})];
if (!group)
group =
ffi::make_object<DeviceHelperSpecializer>(mod, result, names,
target.value(), config);
diff --git a/src/sym/int_set.cc b/src/sym/int_set.cc
index 864416232..4e51eacae 100644
--- a/src/sym/int_set.cc
+++ b/src/sym/int_set.cc
@@ -51,7 +51,10 @@ using prim::IsOne;
using prim::IsZero;
using prim::MakeConst;
-TVM_FFI_STATIC_INIT_BLOCK() { IntervalSetObj::RegisterReflection(); }
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::ObjectDef<IntSetObj>();
+ IntervalSetObj::RegisterReflection();
+}
IntervalSet::IntervalSet(PrimExpr min_value, PrimExpr max_value) :
IntSet(ffi::UnsafeInit{}) {
auto node = ffi::make_object<IntervalSetObj>(std::move(min_value),
std::move(max_value));
diff --git a/src/sym/iter_affine_map.cc b/src/sym/iter_affine_map.cc
index 012c1794c..903c653e4 100644
--- a/src/sym/iter_affine_map.cc
+++ b/src/sym/iter_affine_map.cc
@@ -42,6 +42,7 @@ namespace sym {
using namespace tvm::prim;
TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::ObjectDef<IterMapExprObj>();
IterMarkObj::RegisterReflection();
IterSplitExprObj::RegisterReflection();
IterSumExprObj::RegisterReflection();
diff --git a/src/target/target.cc b/src/target/target.cc
index faa73ff46..cae0f99e0 100644
--- a/src/target/target.cc
+++ b/src/target/target.cc
@@ -54,7 +54,6 @@ class TargetInternal {
ffi::ObjectPtr<TargetObj> n = ffi::make_object<TargetObj>(*target.get());
n->host =
target_host.has_value() ?
ffi::Optional<ffi::ObjectRef>(target_host.value()) : std::nullopt;
- n->str_repr_ = ffi::String();
return Target(ffi::UnsafeInit{}, std::move(n));
}
@@ -94,12 +93,7 @@ static TargetKind GetTargetKind(const ffi::String& name) {
return kind.value();
}
-ffi::String TargetObj::str() const {
- if (str_repr_.empty()) {
- str_repr_ = ffi::json::Stringify(ToConfig());
- }
- return str_repr_;
-}
+ffi::String TargetObj::str() const { return ffi::json::Stringify(ToConfig()); }
/********** Small member methods **********/
diff --git a/src/tirx/ir/layout/tile_core.cc b/src/tirx/ir/layout/tile_core.cc
index 035dc0780..9c4b057dc 100644
--- a/src/tirx/ir/layout/tile_core.cc
+++ b/src/tirx/ir/layout/tile_core.cc
@@ -147,7 +147,10 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
TileLayoutMaybeInplaceM
} // namespace
-TVM_FFI_STATIC_INIT_BLOCK() { ComposeLayoutObj::RegisterReflection(); }
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::ObjectDef<LayoutObj>();
+ ComposeLayoutObj::RegisterReflection();
+}
/**************** Iter ****************/
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/tirx/transform/split_host_device.cc
b/src/tirx/transform/split_host_device.cc
index 121fafe96..ec673a021 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -23,6 +23,7 @@
*/
#include <tvm/backend/cuda/attr.h>
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
@@ -761,7 +762,7 @@ class DeviceKernelMutator : public StmtExprMutator {
}
if (!force_kernel_launch) {
- bool same_target = caller_target->str() == callee_target->str();
+ bool same_target = ffi::StructuralEqual()(caller_target, callee_target);
if (same_target) {
// Calls within the same target may be handled at codegen time
// as internal subroutine calls.
