tilelang-tvm-ir

TileLang TVM IR Handle Conventions

Core Rule

In TVM C++, For, Buffer, Var, SBlock, Stmt, PrimExpr, SeqStmt, etc. are ObjectRef smart handles. ForNode, BufferNode, VarNode, SBlockNode, etc. are raw node structs reached through visitor callbacks, as<TNode>(), operator->, or .get().

When a value needs to cross a function boundary, be stored, be optional, be used as an identity key, or survive beyond a local inspection branch, prefer the handle type over const *Node.

Preferred Patterns

• Function parameters and return values: use handles such as For, Buffer, Var, SBlock, Stmt, or SeqStmt.
• Nullable AST values: use Optional<For>, Optional<SeqStmt>, etc., not const ForNode* = nullptr.
• Identity maps and sets: use handle keys with TVM identity hashing:

using BufferSet = std::unordered_set<Buffer, ObjectPtrHash, ObjectPtrEqual>;
using BufferMap = std::unordered_map<Buffer, Buffer, ObjectPtrHash, ObjectPtrEqual>;
using VarMap = std::unordered_map<Var, PrimExpr, ObjectPtrHash, ObjectPtrEqual>;

• Identity comparisons: use .same_as(other) when comparing two handles.
• Visitor callback node pointers: convert to a handle with GetRef<T>(op) when the value must be retained or passed elsewhere.

Stmt VisitStmt_(const ForNode* op) final {
  For loop = GetRef<For>(op);
  Optional<For> candidate = FindPipelineLoop(loop->body);
  if (candidate.defined() && candidate.value().same_as(loop)) {
    ...
  }
}

• Pattern matching and local mutation may still use node pointers:
  • if (const auto* seq = stmt.as<SeqStmtNode>()) { ... }
  • BufferStoreNode* n = store.CopyOnWrite();
  • visitor overrides such as VisitStmt_(const SeqStmtNode* op)

Keep these raw pointers local to the immediate inspection or mutation site.

TileLang Intrinsic Calls

When adding a new internal TileLang intrinsic, represent it as a registered TileLang Op plus call_intrin, not as call_extern.

• Register the intrinsic in C++ first, usually with a TVM_DLL const Op& declaration in src/op/builtin.h and a matching TIR_DEFINE_TL_BUILTIN(...) registration in src/op/builtin.cc.
• In C++ transforms and lowering passes, build the call with the registered op, for example Call(dtype, tl::my_intrin(), args).
• In Python/TScript helpers or tests, use tirx.call_intrin(...) or T.call_intrin(..., tirx.op.Op.get("tl.my_intrin"), ...).
• Add backend codegen handling by matching the op identity, for example op->op.same_as(tl::my_intrin()), and set any required template/header flags from that intrinsic branch.

Do not use builtin::call_extern, tirx.call_extern, or T.call_extern to smuggle internal TileLang runtime-template helpers through IR. Reserve call_extern for real external ABI calls, user-provided/custom externs, or lowered packed/runtime calls whose contract is genuinely external. Avoid special-casing PrintCallExtern by global symbol name for new TileLang intrinsics; the include and emission logic should live on the intrinsic's call_intrin codegen path.

Avoid

• Passing const ForNode*, const BufferNode*, const VarNode*, or const SBlockNode* between helper functions when a handle exists.
• Storing raw node pointers in std::unordered_map or std::unordered_set for identity tracking.
• Using .get() as a key unless a callee requires a raw TVM node API and the pointer is not retained.
• Comparing handles through .get() == other.get(); prefer .same_as().
• Reconstructing handles from raw pointers repeatedly when a handle is already available.

Common Refactors

// Before
const SeqStmtNode* pipeline_body_seq = nullptr;
pipeline_body_seq = seq_stmt;
ICHECK(pipeline_body_seq != nullptr);

// After
Optional<SeqStmt> pipeline_body_seq;
pipeline_body_seq = GetRef<SeqStmt>(seq_stmt);
ICHECK(pipeline_body_seq.defined());
SeqStmt pipeline_body = pipeline_body_seq.value();

// Before
std::unordered_set<const BufferNode*> seen;
seen.insert(buffer.get());
if (seen.count(read->buffer.get())) { ... }

// After
BufferSet seen;
seen.insert(buffer);
if (seen.count(read->buffer)) { ... }

// Before
std::unordered_set<const VarNode*> vars;
vars.insert(loop->loop_var.get());
bool uses = UsesVar(expr, [&](const VarNode* vn) {
  return vars.count(vn) > 0;
});

// After
VarSet vars;
vars.insert(loop->loop_var);
bool uses = UsesVar(expr, [&](const VarNode* vn) {
  return vars.count(GetRef<Var>(vn)) > 0;
});

Review Checklist

When reviewing TileLang TIR passes, search for:

rg -n "std::unordered_(set|map)<const .*Node \\*|const (For|SeqStmt|SBlock).*Node \\*|\\.get\\(\\) ==|\\.find\\([^\\n]*\\.get\\(\\)|\\.count\\([^\\n]*\\.get\\(\\)|\\.insert\\([^\\n]*\\.get\\(\\)" src/transform

Do not mechanically remove every raw node pointer. Keep visitor signatures, as<TNode>() pattern checks, and CopyOnWrite() mutation pointers. Refactor only the places that store, pass, compare, or key identities through raw pointers.