Skip to content

[FEAT][RUST] Define IR types in Rust - #27

Merged
tqchen merged 1 commit into
apache:mainfrom
Seven-Streams:main-dev/2026-10-21/rust-node-types
Oct 23, 2026
Merged

tqchen merged 1 commit into
apache:mainfrom
Seven-Streams:main-dev/2026-10-21/rust-node-types

Conversation

@Seven-Streams

@Seven-Streams Seven-Streams commented Oct 22, 2026 •

Copy link
Copy Markdown
Member

Summary

Rust could build the IR types C++ defines, but not define a new type of expressions, as
PrimType and TupleType are: a Rust type could not extend ir.Type with fields that C++ and
Python see.

This bumps 3rdparty/tvm-ffi to 8fd9303, the head of apache/tvm-ffi main, which has the merged
apache/tvm-ffi#869 on top of apache/tvm-ffi#870. Its #[rs_class] defines an object type from
its object struct, whose first field base is the parent's object, and registers the reflection
of every other field; def_object_ref! defines the reference to it. This PR also regenerates the
bindings with that stubgen:

  • A complete type that is not final now has a public <Name>Obj::new and implements
    ShallowCopy, which rs_class needs of the parent. This covers ir.Type, ir.Expr,
    ir.Stmt, ir.Attrs and seven more types.
  • The wrapper's new takes a scalar field as its own type and any other as impl Into<T>, as
    the new of an rs_class object does. The hand-written callers drop the .into() they
    applied to an argument, which the new parameter types make ambiguous.
  • Under custom-new, the generated allocator is from_fields instead of
    from_complete_fields ([STUBGEN] Rename the custom-new allocator to from_fields tvm-ffi#870), and the hand-written constructors call it by that
    name.

A Rust type can then extend ir.Type:

use tvm_tirx::ir::{Location, PrimType, Type, TypeObj};
use tvm_tirx::tvm_ffi::{self, def_object_ref, rs_class};

/// The type of a pair of values.
#[rs_class("my_ext.PairType", final)]
pub struct PairTypeNode {
    base: TypeObj,
    pub first: Type,
    pub second: Type,
    /// Not part of structural equality or hashing.
    #[reflect(seq_hash_ignore)]
    pub note: tvm_ffi::String,
}

def_object_ref!(pub PairType, Type, PairTypeNode);

impl PairType {
    pub fn new(first: impl Into<Type>, second: impl Into<Type>) -> Self {
        PairTypeNode::new(TypeObj::new(Location::unknown()), first, second, "").into()
    }
}

let pair = PairType::new(PrimType::int(32), PrimType::float(32));
let ty: Type = pair.into(); // an `ir.Type` that C++ sees as `my_ext.PairType`

The C++ structural equality and hashing, deep copy, repr, the TVMScript printer and Python see
its fields as they see a C++ type's. tests/ir/type.rs checks this end to end:

  • it registers under ir.Type as a TreeNode, with each field's offset and flags;
  • the C ABI getters read its own fields and the inherited loc;
  • the C++ engines compare and hash it, and skip a seq_hash_ignore field;
  • repr and TVMScript print it;
  • C++ deep-copies it, also as the type of a Var inside an Evaluate;
  • the Rust structural engine maps it, a copy-on-write edit keeps the other fields, and it
    converts through Any and the type hierarchy.

@Seven-Streams
Seven-Streams marked this pull request as ready for review October 22, 2026 16:36
@Seven-Streams
Seven-Streams marked this pull request as draft October 22, 2026 18:15
@Seven-Streams
Seven-Streams force-pushed the main-dev/2026-10-21/rust-node-types branch from 349c17b to f958f39 Compare October 23, 2026 00:10
@Seven-Streams
Seven-Streams marked this pull request as ready for review October 23, 2026 00:11
@Seven-Streams
Seven-Streams marked this pull request as draft October 23, 2026 00:12
@Seven-Streams
Seven-Streams force-pushed the main-dev/2026-10-21/rust-node-types branch from f958f39 to 54eade2 Compare October 23, 2026 00:19
@Seven-Streams
Seven-Streams marked this pull request as ready for review October 23, 2026 00:19
@Seven-Streams
Seven-Streams force-pushed the main-dev/2026-10-21/rust-node-types branch from 54eade2 to 786f7f4 Compare October 23, 2026 01:17
@tqchen
tqchen merged commit 503e143 into apache:main Oct 23, 2026
6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants