tilelang.testing.ir¶

Small, exact queries for compiler tests, without running any passes.

Counts describe unique static IR nodes within each visited function, following TVM’s post-order visitor. They do not describe dynamic execution counts or order.

Functions¶

collect_nodes(root, node_type, *[, function])

Collect nodes of a TVM object type from bodies or an expression.

collect_calls(root, *, op[, function])

Collect calls to one exact registered Op; never match name fragments.

assert_call_count(root, *, op, count[, function])

Assert a static call count, showing function ownership on failure.

collect_extern_calls(root, *, symbol[, function])

Collect tirx.call_extern calls with an exact StringImm symbol argument.

Module Contents¶

tilelang.testing.ir.collect_nodes(root, node_type, *, function=None)¶

Collect nodes of a TVM object type from bodies or an expression.

A module query visits all PrimFunc bodies unless function selects one by global name. Other module functions are skipped. Function signatures and module metadata are not traversed. Shared nodes are visited once per body.

Parameters:

function (str | None)

Return type:

list

tilelang.testing.ir.collect_calls(root, *, op, function=None)¶

Collect calls to one exact registered Op; never match name fragments.

GlobalVar and other non-Op callees do not match. An unregistered name raises the Op registry’s error instead of silently producing an empty result.

Parameters:

  • op (str | tilelang.tvm.ir.Op)

  • function (str | None)

Return type:

list[tvm.tirx.Call]

tilelang.testing.ir.assert_call_count(root, *, op, count, function=None)¶

Assert a static call count, showing function ownership on failure.

Parameters:

  • op (str | tilelang.tvm.ir.Op)

  • count (int)

  • function (str | None)

Return type:

None

tilelang.testing.ir.collect_extern_calls(root, *, symbol, function=None)¶

Collect tirx.call_extern calls with an exact StringImm symbol argument.

Other call operators, missing arguments and non-string symbols do not match. This function does not interpret prefixes, templates or regular expressions.

Parameters:

  • symbol (str)

  • function (str | None)

Return type:

list[tvm.tirx.Call]