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 72c99622 [REFACTOR] Require an explicit type in the Var constructor
(#19)
72c99622 is described below
commit 72c996225bec47e58d70c6c723d962feb92c687e
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 21 20:02:49 2026 -0400
[REFACTOR] Require an explicit type in the Var constructor (#19)
Require a concrete `Type` in the canonical C++ `Var(name, ty, loc)`
constructor and its FFI binding, storing it directly in the inherited
expression type field. Python and Rust retain their optional-type APIs
by translating an omitted type to `MissingType()` before calling the
typed constructor.
---
include/tvm/ir/expr.h | 3 +--
python/tvm_tirx/ir/expr.py | 8 +++++---
rust/src/ir/expr.rs | 6 +++++-
src/ir/expr.cc | 11 ++++-------
4 files changed, 15 insertions(+), 13 deletions(-)
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index 879f8168..e0c2a109 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -571,8 +571,7 @@ class VarNode : public ExprNode {
/*! \brief Managed reference to VarNode. */
class Var : public Expr {
public:
- TVM_TIRX_DLL explicit Var(ffi::String name, ffi::Optional<Type>
ty_annotation,
- Location loc = UnknownLoc());
+ TVM_TIRX_DLL explicit Var(ffi::String name, Type ty, Location loc =
UnknownLoc());
/*! \brief Return a fresh ordinary Var with the same type and a new name. */
TVM_TIRX_DLL Var CopyWithName(const ffi::String& name) const;
diff --git a/python/tvm_tirx/ir/expr.py b/python/tvm_tirx/ir/expr.py
index 84f5c762..b7467b02 100644
--- a/python/tvm_tirx/ir/expr.py
+++ b/python/tvm_tirx/ir/expr.py
@@ -843,11 +843,13 @@ class Var(_CallableExprWithOp):
raise TypeError("name must be a str")
# pylint: disable=import-outside-toplevel
- from .type import PrimType, PtrType, Type
+ from .type import MissingType, PrimType, PtrType, Type
- if isinstance(ty, str):
+ if ty is None:
+ ty = MissingType()
+ elif isinstance(ty, str):
ty = PtrType(PrimType("void")) if ty == "handle" else PrimType(ty)
- elif ty is not None:
+ else:
ty = tvm_ffi.convert(ty)
if not isinstance(ty, Type):
raise TypeError("ty must be a Type or primitive dtype string")
diff --git a/rust/src/ir/expr.rs b/rust/src/ir/expr.rs
index e36be9bd..69fc7999 100644
--- a/rust/src/ir/expr.rs
+++ b/rust/src/ir/expr.rs
@@ -1116,8 +1116,12 @@ tvm_ffi::impl_object_upcast!(Var => Expr);
impl Var {
/// The variable `name`, of type `ty` or a missing type when `None`
(`ir.Var`).
pub fn new(name: &str, ty: Option<&Type>, loc: impl Into<Location>) ->
Result<Self> {
+ let ty = match ty {
+ Some(ty) => ty.clone(),
+ None => crate::ir::r#type::MissingType::new()?.into(),
+ };
tvm_ffi::cached_global_func!("ir.Var")
- .call_tuple((String::from(name), ty.cloned(), &loc.into()))?
+ .call_tuple((String::from(name), ty, &loc.into()))?
.try_into()
}
}
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index 3ac973f7..fbf90fa5 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -1182,13 +1182,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
// Var
-Var::Var(ffi::String name, ffi::Optional<Type> ty_annotation, Location loc)
- : Expr(ffi::UnsafeInit{}) {
+Var::Var(ffi::String name, Type ty, Location loc) : Expr(ffi::UnsafeInit{}) {
ffi::ObjectPtr<VarNode> n = ffi::make_object<VarNode>();
n->name = std::move(name);
- if (ty_annotation.has_value()) {
- n->ty = ty_annotation.value();
- }
+ n->ty = std::move(ty);
n->loc = loc;
data_ = std::move(n);
}
@@ -1202,8 +1199,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(ffi::type_attr::kStructuralMaybeInplaceMutate,
ffi::FStructuralMutate::FromNative<&VarMaybeInplaceMutate>());
- refl::GlobalDef().def("ir.Var", [](ffi::String name, ffi::Optional<Type>
ty_annotation,
- Location loc) { return Var(name,
ty_annotation, loc); });
+ refl::GlobalDef().def("ir.Var",
+ [](ffi::String name, Type ty, Location loc) { return
Var(name, ty, loc); });
}
Var Var::CopyWithName(const ffi::String& name) const {
