Working with IR in C++#

TIRx is implemented in C++, and a C++ function can read, build, and rewrite its IR directly. The function can receive a program written in TVMScript, inspect or rewrite it, and return the result, with no conversion in between, which lets us move a hot analysis or a large transformation to C++. This guide covers:

TVM-FFI object system: how IR nodes are represented and used in C++.
Structural utilities: walk, map, and compare programs.
Visitor and mutator classes: one hook per node type, for larger passes.

Note

The runnable code in this guide is under examples/ir_guides in the repository.

Setting Up#

The examples are packaged as tirx_ir_guide_examples under examples/ir_guides. Installing the package builds the C++ source, cpp/working_with_ir.cc, into a library; CMakeLists.txt and pyproject.toml in the same directory show how. Install TIRx from the same source tree to keep the IR APIs aligned. We install without build isolation to build against that installation instead of fetching one from PyPI:

pip install scikit-build-core
pip install --no-build-isolation -e examples/ir_guides

A TVM_FFI_STATIC_INIT_BLOCK at the end of the source registers each example as a global function under tirx_ir_guide_examples.cpp:

refl::GlobalDef()
    .def("tirx_ir_guide_examples.cpp.construct", Construct)

The package’s tirx_ir_guide_examples.cpp module loads the library and exposes these functions, so Python calls them by name. The driver working_with_ir_in_cpp.py calls each example this way and prints the result.

The object system examples take no arguments: each builds its own IR in C++. The later sections receive the add_one function from Quick Start, written in TVMScript:

from tvm_tirx.script import tirx as T

n = T.dynamic("n")


@T.function
def add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
    for i in range(n):
        B[i] = A[i] + 1

import tirx_ir_guide_examples.cpp

tirx_ir_guide_examples.cpp.construct()
tirx_ir_guide_examples.cpp.tensors_written(add_one)

The headers relevant to this guide sit in five folders under tvm/; see Header Organization for the full list:

Folder	Contents	Included by our examples
`tvm/ffi/`	From TVM-FFI: `Object` and `ObjectRef`, containers, functions, and the structural utilities.	`extra/structural_visit.h`, `extra/structural_mutate.h`, `extra/structural_equal.h`, `extra/structural_hash.h`, `reflection/registry.h`
`tvm/ir/`	Core IR nodes, such as `Expr`, `Var`, `Stmt`, `For`, and `TensorStore`, and the expression functors in `expr_functor.h`.	`expr_functor.h`
`tvm/sym/`	Symbolic analysis, such as `sym::Analyzer` and `sym::ConstraintContext`.	None
`tvm/tirx/`	Tensor programs: `tirx::Function`, the statement visitors in `stmt_functor.h`, and the constraint-context variants.	`function.h`, `stmt_functor.h`
`tvm/script/`	The IR builder and the TVMScript printer.	None

IR classes live in the tvm namespace, primitive expressions such as tvm::prim::Add in tvm::prim, TIRx-specific ones such as tvm::tirx::Function in tvm::tirx, and the TVM-FFI classes in tvm::ffi. The examples start with using namespace tvm;, so these appear as Var, prim::Add, tirx::Function, and ffi::StructuralWalk.

TVM-FFI Object System#

Every TIRx IR node is a TVM-FFI object. TIRx is designed for dynamic extension: a library can add a new node type, for example a new statement for a hardware feature, and existing code keeps working on programs that contain it, which a fixed class hierarchy compiled into one library cannot offer. Generic utilities such as walking, rewriting, comparing, and printing work on every node through its registered fields.

TVM-FFI objects also work across languages. Python, C++, and Rust share the same native, immutable, reference-counted objects, so a program moves between them without conversion.

In C++, TVM-FFI objects are ordinary C++ classes, used through a few key constructs:

Construct	Purpose
Constructors, `PrimExpr` operators	Build expressions with ordinary C++ operators
`FooObj` / `Foo`	Underlying object and its named managed reference: read fields with `->`, copy cheaply
`CopyOnWrite()`	Change a field, copying the object only if it is shared
Implicit upcast	Pass a node wherever a parent type is expected
`as<T>()` / `as_or_throw<T>()`	Convert to a child type, with a checked result
Typed views (`PrimExpr`, `IntExpr`, `PrimVar`)	Require an expression of a given type, checked on conversion
`ExprFunctor` / `StmtFunctor`	Process each node type in its own method

The subsections below introduce each construct.

Constructing Objects#

We build IR with constructors and ordinary C++ operators. Writing x * 2 + 1 as nested constructor calls would bury the expression, so arithmetic operators on PrimExpr build the nodes for us. Let’s build it:

PrimExpr XTimes2Plus1() {
  PrimVar x("x", PrimType::Int(32));
  return x * 2 + 1;
}

ffi::String Construct() { return Print(XTimes2Plus1()); }

The output is:

x * 2 + 1

PrimVar creates an integer variable x, a variable known to have a primitive type. * and + turn the literals 2 and 1 into constants of the type of x, fold constants, and otherwise call the node constructors, such as prim::Add(a, b). Writing a node to a stream prints it as TVMScript; Print is our small helper that does so.

We now have an expression; next, let’s see what kind of C++ value it is.

Objects and Managed References#

Every node type comes as a pair: an underlying object and a named managed reference. The top of x * 2 + 1 is an addition, so let’s view the expression as an addition, with a cast the next subsection explains, and read its two operands:

ffi::Array<ffi::String> Objects() {
  prim::Add add = XTimes2Plus1().as_or_throw<prim::Add>();
  prim::Add alias = add;
  return {"a = " + Print(add->a), "b = " + Print(add->b),
          std::string("same object: ") + (alias.same_as(add) ? "true" : "false")};
}

The output is:

a = x * 2
b = 1
same object: true

alias is the same object as add: copying a managed reference copies only the pointer, and same_as compares identity.

AddObj is the underlying object: the class that holds the fields. Add is a named managed reference to it: an owning, reference-counted pointer, like std::shared_ptr<AddObj> except that it is never null, so copying it is cheap and -> is always safe. Functions take and return managed references, and -> on a managed reference reaches the object’s fields: add->a.

We define a separate named class Add instead of using a plain pointer to AddObj because the named class carries the interface: constructors, methods, and the type information that lets it pass through TVM-FFI functions. It can also impose additional constraints on the underlying object that the object alone does not express.

Nodes are immutable once shared. To change a field, CopyOnWrite() returns a mutable pointer to the object, copying the object first if anyone else holds it:

ffi::Array<ffi::String> CopyOnWrite() {
  prim::Add add = XTimes2Plus1().as_or_throw<prim::Add>();
  prim::Add changed = add;
  changed.CopyOnWrite()->b = 3;
  return {"add = " + Print(add), "changed = " + Print(changed)};
}

The output is:

add = x * 2 + 1
changed = x * 2 + 3

changed shared its object with add, so CopyOnWrite() copied it before the write, and add is untouched.

We reached Add with a cast; next, let’s look at casts on their own.

Casts#

Casts move a managed reference up and down the class hierarchy. A function that takes any expression wants an Expr, while reading the operands needs an Add. Converting up is implicit, through C++ inheritance. Converting down checks the node’s actual type, with ObjectRef::as and its relatives in one of three forms:

Call	Returns
`value.as<AddObj>()`	A non-owning `const AddObj*`, valid while a reference keeps the object alive, or `nullptr` if it is not an addition.
`value.as<Add>()`	A `std::optional<Add>`.
`value.as_or_throw<Add>()`	The `Add`, or raises an error.

ffi::Array<ffi::String> Casts() {
  Expr value = XTimes2Plus1();
  const prim::AddObj* node = value.as<prim::AddObj>();
  std::optional<prim::Add> add = value.as<prim::Add>();
  std::optional<prim::Mul> mul = value.as<prim::Mul>();
  ffi::Array<ffi::String> lines = {
      std::string("as<AddObj>: ") + (node != nullptr ? "pointer" : "nullptr"),
      std::string("as<Add>: ") + (add.has_value() ? "Add" : "nullopt"),
      std::string("as<Mul>: ") + (mul.has_value() ? "Mul" : "nullopt")};
  try {
    value.as_or_throw<prim::Mul>();
  } catch (const ffi::Error& err) {
    lines.push_back("as_or_throw<Mul>: " + std::string(err.message()));
  }
  return lines;
}

The output is:

as<AddObj>: pointer
as<Add>: Add
as<Mul>: nullopt
as_or_throw<Mul>: Cannot treat type `prim.Add` as type `prim.Mul`

value is an Expr, but the node behind it is still an addition, so the casts to AddObj and Add succeed and the cast to Mul finds nothing. as<AddObj>() works like dynamic_cast<const AddObj*>, and as<Add>() and as_or_throw<Add>() do the same check but return a managed reference. as_or_throw turns the same mismatch into an error that names both type keys: the actual prim.Add and the requested prim.Mul.

Casts check the class of the object; next, let’s look at references that also check the type of the expression.

Typed Views#

A typed view constrains an expression by its type, not by its object class. Arithmetic needs expressions whose result is a primitive value, but an addition, a variable, and a call can all produce one, so no single object class captures “primitive”. PrimExpr is a named managed reference that accepts any Expr whose ty is a PrimType, and has no object class of its own.

View	Accepts
`PrimExpr`	Any expression whose `ty` is a `PrimType`
`IntExpr`	A scalar integer expression, signed or unsigned
`PrimVar`	A `Var` whose `ty` is a `PrimType`

The same casts check the view. Let’s try two variables, an integer and a float:

ffi::Array<ffi::String> TypedViews() {
  Var i("i", PrimType::Int(32));
  Var f("f", PrimType::Float(32));
  ffi::Array<ffi::String> lines = {
      std::string("i as PrimExpr: ") + (i.as<PrimExpr>().has_value() ? "yes" : "no"),
      std::string("i as IntExpr: ") + (i.as<IntExpr>().has_value() ? "yes" : "no"),
      std::string("f as IntExpr: ") + (f.as<IntExpr>().has_value() ? "yes" : "no")};
  try {
    f.as_or_throw<IntExpr>();
  } catch (const ffi::Error& err) {
    lines.push_back("as_or_throw<IntExpr>: " + std::string(err.message()));
  }
  return lines;
}

The output is:

i as PrimExpr: yes
i as IntExpr: yes
f as IntExpr: no
as_or_throw<IntExpr>: Cannot treat type `ir.Var[ty=float32]` as type `ir.IntExpr`

Both variables are VarObj objects; only their ty differs, and the views check it. The error shows the type in brackets, ir.Var[ty=float32]. This is why x is a PrimVar: the operators take PrimExpr, and a plain Var would first need a checked cast.

Note

A view is not a subclass of the node’s own reference class: PrimVar derives from PrimExpr and converts to Var through a conversion operator.

A cast or a view looks at one node at a time; to process every node type in its own method, we dispatch.

Dispatching on Node Types#

A functor picks one of its own methods by the node’s runtime type, as in the visitor pattern. ExprFunctor declares a virtual Dispatch_ overload for each expression class, and a subclass overrides the ones it cares about. Dispatch looks up the node’s type in a table and calls the matching overload, so the nodes need no accept() method. Let’s print an expression with explicit parentheses:

class Describe : public ExprFunctor<std::string(const Expr&)> {
 public:
  using ExprFunctor::Dispatch_;

  std::string Dispatch_(const prim::AddObj* add) override {
    return "(" + Dispatch(add->a) + " + " + Dispatch(add->b) + ")";
  }
  std::string Dispatch_(const prim::MulObj* mul) override {
    return "(" + Dispatch(mul->a) + " * " + Dispatch(mul->b) + ")";
  }
  std::string Dispatch_(const VarObj* var) override { return var->name; }
  std::string Dispatch_(const IntImmObj* imm) override { return Print(ffi::GetRef<IntImm>(imm)); }
  std::string DispatchDefault_(const ffi::Object*) override { return "<other>"; }
};

We call it on x * 2 + 1 and on x - 1:

ffi::Array<ffi::String> Dispatching() {
  PrimVar x("x", PrimType::Int(32));
  Describe describe;
  return {describe(XTimes2Plus1()), describe(x - 1)};
}

The output is:

((x * 2) + 1)
<other>

The template argument chooses the signature, here std::string from an Expr. A functor does not visit children on its own: the Add and Mul overloads recurse by calling Dispatch on the operands, and the Var and IntImm overloads end the recursion.

Unhandled node types fall back to a default method. They go to DispatchDefault_, which raises an error unless we override it. Our override returns <other>, so x - 1 passes through instead of failing. tirx::StmtFunctor does the same for statements.

Note

The using ExprFunctor::Dispatch_; line keeps the overloads we did not override visible, as usual for overloaded virtual functions in C++.

Dispatch looks at one node and recurses by hand; the structural utilities below walk whole programs for us.

Structural Utilities in C++#

The structural utilities are part of TVM-FFI and work on every node through its registered fields: ffi::StructuralWalk, ffi::StructuralMap, ffi::StructuralEqual, and ffi::StructuralHash. Each piece takes a C++ form:

Piece	C++ form
Nodes a callback receives	Selected by the type of the callback’s first parameter
Walk callback result	`ffi::Expected<WalkResult>`
Traversal order	A required template argument, `WalkOrder::kPreOrder` or `kPostOrder`
Keep a node in a map	Return `ffi::Unchanged()` or the node
Result of a map	An `ffi::Any`, cast back with `cast<T>()`
Definition region kind	A trailing `TVMFFIDefRegionKind` parameter

Structural Walk#

A walk calls our callback on each node of the type it asks for. To find the tensors that add_one writes, the callback takes a const TensorStore&:

ffi::Array<ffi::String> TensorsWritten(const tirx::Function& func) {
  ffi::Array<ffi::String> writes;
  ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(
      func, [&](const TensorStore& store) -> ffi::Expected<ffi::WalkResult> {
        writes.push_back(store->dest.as<VarObj>()->name);
        return ffi::WalkResult::Advance();
      });
  return writes;
}

The output is:

['B']

Only the store reaches the callback, because its parameter is const TensorStore&.

A walk callback returns an ffi::Expected<ffi::WalkResult>. WalkResult::Advance() continues the walk, WalkResult::Skip() skips the node’s children in pre-order, and WalkResult::Interrupt(...) stops the walk.

To handle several node types in one walk, we pass one callback per type as extra arguments. They are tried in declaration order, and only the first one whose parameter type matches a node runs; a node that matches none is still traversed but not reported. This walk records the loops and stores in pre-order:

ffi::Array<ffi::String> TraceStatements(const tirx::Function& func) {
  ffi::Array<ffi::String> trace;
  ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
      func,
      [&](const For& loop) -> ffi::Expected<ffi::WalkResult> {
        trace.push_back("loop " + loop->var.as<VarObj>()->name);
        return ffi::WalkResult::Advance();
      },
      [&](const TensorStore& store) -> ffi::Expected<ffi::WalkResult> {
        trace.push_back("store " + store->dest.as<VarObj>()->name);
        return ffi::WalkResult::Advance();
      });
  return trace;
}

The output is:

['loop i', 'store B']

The loop and the store each reach their own callback, and pre-order reports the loop before the store inside it. Matching follows the class hierarchy, so a callback that takes a base class, such as const Stmt&, would also take loops and stores that no earlier callback claims.

To stop at the first match and return it, a callback wraps the node in an ffi::VisitInterrupt. StructuralWalk returns that interrupt, or std::nullopt if the walk finished without one:

ffi::Optional<TensorStore> FindWrite(const tirx::Function& func) {
  ffi::Optional<ffi::VisitInterrupt> found = ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
      func, [](const TensorStore& store) -> ffi::Expected<ffi::WalkResult> {
        return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(store));
      });
  if (!found.has_value()) return std::nullopt;
  return found.value()->value.cast<TensorStore>();
}

The output is:

The interrupt’s value field carries the payload as an ffi::Any, and cast<TensorStore>() converts it back to a store.

Definitions and Uses#

A callback receives the definition region kind of each site when it takes a TVMFFIDefRegionKind as its last parameter: kTVMFFIDefRegionKindNone for a use, and kTVMFFIDefRegionKindSimple and kTVMFFIDefRegionKindPattern for the two kinds of definition site. Let’s list the definition sites in add_one:

ffi::Array<ffi::String> Definitions(const tirx::Function& func) {
  ffi::Array<ffi::String> sites;
  ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
      func, [&](const Var& var, TVMFFIDefRegionKind kind) -> ffi::Expected<ffi::WalkResult> {
        if (kind == kTVMFFIDefRegionKindSimple) {
          sites.push_back(var->name + " simple");
        } else if (kind == kTVMFFIDefRegionKindPattern) {
          sites.push_back(var->name + " pattern");
        }
        return ffi::WalkResult::Advance();
      });
  return sites;
}

The output is:

A pattern
n pattern
B pattern
n pattern
i simple

The parameters define A, B, and n by pattern, and the loop defines i.

Structural Map#

ffi::StructuralMap takes the same callbacks, but each one returns a replacement for its node. Let’s specialize add_one by replacing n with 128, passing n in from Python:

tirx::Function Specialize(const tirx::Function& func, const Var& n) {
  ffi::Any result =
      ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(func, [&](const Var& var) -> ffi::Any {
        if (var.same_as(n)) {
          return IntImm(var->ty.as_or_throw<PrimType>(), 128);
        }
        return ffi::Unchanged();
      });
  return result.cast<tirx::Function>();
}

The output is:

from __future__ import annotations

# from tvm_tirx.script import tirx as T

@T.function
def add_one(A: T.Tensor((128,), T.float32), B: T.Tensor((128,), T.float32)):
    for i in range(128):
        B[i] = A[i] + T.float32(1.0)

For every other variable, the callback returns ffi::Unchanged(), which keeps the node as it is. The constant takes its type from the variable, var->ty. The mapper keeps definitions and uses consistent, so the new A and B and every use of them carry the specialized shape.

Equality and Hashing#

ffi::StructuralEqual::Equal and ffi::StructuralHash::Hash compare and hash by structure. Let’s compare two add_one functions built separately in Python by make_add(1), the helper from the structural utilities guide:

ffi::Array<bool> Compare(const tirx::Function& first, const tirx::Function& second) {
  return {
      first.same_as(second),
      ffi::StructuralEqual::Equal(first, second),
      ffi::StructuralHash::Hash(first) == ffi::StructuralHash::Hash(second),
  };
}

The output is:

[False, True, True]

The functions are different objects, yet structurally equal with equal hashes. Equal and Hash also take an optional map_free_vars argument, which defaults to false.

Visit and Mutate#

When a callback must decide which children to visit, we use ffi::StructuralVisit and ffi::StructuralMutate. Their callbacks also receive the visitor or mutator, such as ffi::StructuralVisitorObj*, and descend into a node’s children with DefaultVisitExpected or DefaultMutateExpected.

Visitor and Mutator Classes#

Walk and map cover most tasks. A pass that handles many node types, however, reads better as a class with one method per type:

Tool	Covers	Use it for
Structural utilities	Any TVM-FFI object	A few node types, written as callbacks
`tirx::StmtExprVisitor`, `tirx::StmtExprMutator`	TIRx statements and expressions (`tvm/tirx/stmt_functor.h`)	Passes that customize many node types
`ExprVisitor`, `ExprMutator`	Expressions alone (`tvm/ir/expr_functor.h`)	Expression analyses and rewrites
`ExprFunctor`, `tirx::StmtFunctor`	Dispatch on the node type with a signature we choose, such as `int(const Expr&)`, without traversing children	Computations where each hook decides which children to recurse into

Each class declares a virtual hook for every node type, named Visit_ for a visitor and Mutate_ for a mutator, and a subclass overrides only the types it cares about. Unlike the functors, these classes traverse children for us.

These classes are built on the structural utilities. A node type without a hook, such as a function, falls back to its structural traversal, which hands each child back to our hooks.

Note

ExprVisitor ultimately derives from ffi::StructuralVisitorObj, and ExprMutator from the mutator base behind StructuralMap. Any hook can read the definition region kind with def_region_kind().

Let’s record the loops around each write:

class LoopNestVisitor : public tirx::StmtExprVisitor {
 public:
  using tirx::StmtExprVisitor::Visit_;

  ffi::Optional<ffi::VisitInterrupt> Visit_(const ForObj* loop) override {
    loops_.push_back(loop->var.as<VarObj>()->name);
    ffi::Optional<ffi::VisitInterrupt> result = tirx::StmtExprVisitor::Visit_(loop);
    loops_.pop_back();
    return result;
  }

  ffi::Optional<ffi::VisitInterrupt> Visit_(const TensorStoreObj* store) override {
    std::string entry = std::string(store->dest.as<VarObj>()->name) + " in";
    for (const std::string& loop : loops_) entry += " " + loop;
    writes.push_back(entry);
    return tirx::StmtExprVisitor::Visit_(store);
  }

  ffi::Array<ffi::String> writes;

 private:
  std::vector<std::string> loops_;
};

ffi::Array<ffi::String> LoopsAroundWrites(const tirx::Function& func) {
  auto visitor = ffi::make_object<LoopNestVisitor>();
  visitor->Visit(func);
  return visitor->writes;
}

The output is:

['B in i']

Our For hook pushes the loop variable, calls the base class hook to visit the loop’s bounds and body, then pops the variable. Calling the base hook visits the children with their usual traversal, and returning its result passes an interrupt from a child up to the caller. Visitors are objects, so we allocate them with ffi::make_object.

Note

The using tirx::StmtExprVisitor::Visit_; line keeps the hooks we did not override visible, as usual for overloaded virtual functions in C++.

A mutator works the same way, with hooks that return the replacement node. Let’s reverse the loop by replacing each use of i with n - 1 - i, built with arithmetic operators as before:

class ReverseLoopMutator : public tirx::StmtExprMutator {
 public:
  using tirx::StmtExprMutator::Mutate_;

  UnchangedOr<Stmt> Mutate_(const ForObj* loop, InplaceMode mode) override {
    reversed_.emplace(loop->var.as<VarObj>(), loop->end - 1 - loop->var);
    return tirx::StmtExprMutator::Mutate_(loop, mode);
  }

  UnchangedOr<Expr> Mutate_(const VarObj* var, InplaceMode mode) override {
    auto it = reversed_.find(var);
    if (it != reversed_.end()) return it->second;
    return tirx::StmtExprMutator::Mutate_(var, mode);
  }

 private:
  std::unordered_map<const VarObj*, PrimExpr> reversed_;
};

tirx::Function ReverseLoop(const tirx::Function& func) {
  auto mutator = ffi::make_object<ReverseLoopMutator>();
  return mutator->Mutate(func).ValueOrUnchanged(func).as_or_throw<tirx::Function>();
}

The output is:

from __future__ import annotations

# from tvm_tirx.script import tirx as T

@T.function
def add_one[n](A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
    for i in range(n):
        B[n - 1 - i] = A[n - 1 - i] + T.float32(1.0)

A mutator hook returns an UnchangedOr<T>: either a replacement node or ffi::Unchanged(), which keeps the original and lets the parents above it stay unchanged too. Our For hook records the replacement before calling the base class hook, so it is ready before any use of i is reached.

Note

StmtExprMutator’s For hook rewrites the loop’s bounds and body but not the loop variable itself, so our Var hook sees only the uses of i, and the loop keeps its variable. The InplaceMode argument tells a hook whether it may update a node that no one else references in place; our hooks forward it to the base class unchanged.

Constraint Contexts#

TIRx adds variants of the structural utilities and the visitor and mutator classes that also pass the facts valid at each node, such as loop ranges and branch conditions:

Utility	Works like	Facts reach the callback as
`tirx::StructuralWalkWithConstraintContext`	`ffi::StructuralWalk`	A `const sym::ConstraintContext&` after the node
`tirx::StructuralMapWithConstraintContext`	`ffi::StructuralMap`	A `const sym::ConstraintContext&` after the node
`tirx::StructuralVisitWithConstraintContext`	`ffi::StructuralVisit`	The `constraints` member of the visitor the callback receives
`tirx::StructuralMutateWithConstraintContext`	`ffi::StructuralMutate`	The `constraints` member of the mutator the callback receives
`tirx::StmtExprVisitorWithConstraintContext`, `tirx::StmtExprMutatorWithConstraintContext`	`tirx::StmtExprVisitor`, `tirx::StmtExprMutator`	The inherited `constraints` member

The functions are in tvm/tirx/structural_visit_with_constraint_context.h and tvm/tirx/structural_mutate_with_constraint_context.h, and the classes in tvm/tirx/stmt_functor_with_constraint_context.h.

Across these tools, IR is immutable, rewrites return new programs that share unchanged parts, and every utility respects definitions and uses. Because C++ and Python share the same objects, a pass can move between the two languages one function at a time.

Quick Reference#

Construct	What it does
`PrimVar x("x", PrimType::Int(32))`, `x * 2 + 1`	Creates a node; builds an expression with operators
`add->a`, `loop->var`	Reads a node’s fields
`value.as<AddObj>()`, `value.as<Add>()`, `value.as_or_throw<Add>()`	Checks the node type and downcasts
`PrimExpr`, `IntExpr`, `PrimVar`	Typed views that check an expression’s `ty`
`ref.CopyOnWrite()->field = value`	Changes a field, copying the object first if it is shared
`ExprFunctor`, `tirx::StmtFunctor`	One `Dispatch_` overload per node type, with explicit recursion
`ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(func, ...)`	Walks with typed callbacks that return `ffi::Expected<ffi::WalkResult>`
`ffi::StructuralMap<order>(func, ...)`, `ffi::Unchanged()`	Rebuilds with replacements and returns an `ffi::Any`
A trailing `TVMFFIDefRegionKind` parameter	Receives the definition region kind
`ffi::StructuralEqual::Equal`, `ffi::StructuralHash::Hash`	Compares and hashes by structure
`tirx::StmtExprVisitor`, `tirx::StmtExprMutator`, `ExprVisitor`, `ExprMutator`	One `Visit_` or `Mutate_` hook per node type
`TVM_FFI_STATIC_INIT_BLOCK`, `refl::GlobalDef().def(...)`	Registers C++ functions as global functions

Working with IR in C++

Contents

Working with IR in C++#

Setting Up#

TVM-FFI Object System#

Constructing Objects#

Objects and Managed References#

Casts#

Typed Views#

Dispatching on Node Types#

Structural Utilities in C++#

Structural Walk#

Definitions and Uses#

Structural Map#

Equality and Hashing#

Visit and Mutate#

Visitor and Mutator Classes#

Constraint Contexts#

Quick Reference#