Processing MLIR with ZiriumDRAFT

A small question about an MLIR file can require a surprising amount of machinery to answer. What uses this value? Which operations contribute to a particular return value? Did a compiler pass change the dataflow, or merely rename a few values? Text tools are convenient, but following those relationships requires more than matching operation names. Linking LLVM and MLIR into a small utility or Python package is a much larger commitment.

I built Zirium for these kinds of questions. It is a Rust library for parsing textual MLIR without linking LLVM, with typed Python bindings and a command-line tool. The library retains both the source text and a structured representation that tools can inspect and edit.

MLIR is a compiler infrastructure in which operations from different dialects can appear in the same program. A dialect defines operations for a particular purpose, such as arithmetic or tensor computations. Zirium understands generic operation syntax and selected dialect-specific forms. It is early software, with incomplete dialect coverage and an API that may change without a migration path.

The command-line tool is a convenient way to explore the library’s query support. Its query language borrows the pipeline model from jq, with stages that can follow the relationships between MLIR operations.

A five-operation program#

Save this small example as arithmetic.mlir:

module {
  %lhs = arith.constant 6 : i32
  %rhs = arith.constant 7 : i32
  %sum = arith.addi %lhs, %rhs : i32
  %product = "arith.muli"(%sum, %rhs)
    {analysis.tag = "old"} : (i32, i32) -> i32
}

The two constants feed an addition, and the multiplication uses both the sum and the second constant. The multiplication is written in MLIR’s generic quoted syntax; the other operations use custom assembly forms.

Zirium queries start with an implicit stream containing every operation. count reduces that stream to a number:

$ zirium 'count' arithmetic.mlir
5

The module counts too. To see what the five operations are, project their names and tally them:

$ zirium 'names | tally | json' arithmetic.mlir
{
  "arith.addi": 1,
  "arith.constant": 2,
  "arith.muli": 1,
  "builtin.module": 1
}

A stage can change the kind of value flowing through the query. names turns operations into strings. tally turns those strings into a map of counts. json explicitly requests json output, though maps already print as json by default.

Following values through the program#

MLIR uses static single assignment (SSA) values. Each value has one definition, so a tool can follow a use of %sum back to the operation that produced it.

Filtering keeps the operations that match a predicate:

zirium 'filter(op("arith.addi"))' arithmetic.mlir

The result is the selected addition, printed as MLIR inside enough of its enclosing module to make its location clear. Predicates can also be combined. This selects additions and multiplications without an analysis.tag attribute:

filter(
  (op("arith.addi") or op("arith.muli"))
  and not has_attr("analysis.tag")
)

Now follow the result of the addition to its direct users:

$ zirium 'filter(op("arith.addi")) | users | names' arithmetic.mlir
arith.muli

Going the other way, defs follows operands to the operations that define them:

$ zirium 'filter(op("arith.muli")) | defs | names' arithmetic.mlir
arith.addi
arith.constant

The multiplication uses the result of the addition and the %rhs constant. defs(0) would follow only its first operand. Similarly, users(0) follows uses of an operation’s first result, which is useful for operations with several results.

backward_slice (also available as slice) follows definitions transitively backward through the graph and includes the starting operation. In the other direction, forward_slice follows uses transitively forward through the graph.

$ zirium 'filter(op("arith.addi")) | forward_slice | names' arithmetic.mlir
arith.addi
arith.muli

Both directional slice stages follow explicit SSA edges. They do not infer control-flow, symbol, region, memory, or effect dependencies.

The illustration below follows the multiplication back to its definitions, then counts their operation names. Move the slider to see what each stage contributes:

1module {
2  %lhs = arith.constant 6 : i32
3  %rhs = arith.constant 7 : i32
4  %sum = arith.addi %lhs, %rhs : i32
5  %product = "arith.muli"(%sum, %rhs)
6    {analysis.tag = "old"} : (i32, i32) -> i32
7}

To follow definitions beyond the immediate operands shown in the illustration, use backward_slice. The traversal stops at block arguments, such as function parameters, and operations without operands:

$ zirium \
    'filter(op("arith.muli")) | backward_slice | names | tally | json' \
    arithmetic.mlir
{
  "arith.addi": 1,
  "arith.constant": 2,
  "arith.muli": 1
}

In a large compiler dump, the same query gives a compact view of the computation feeding one operation. For a function with several returned tensors, this query starts at the operation defining the first returned value:

filter(op("func.return")) | defs(0) | backward_slice

This traces explicit SSA dependencies. At an operation with several results, backward_slice follows all operands; it does not infer which operands contribute to the selected result. There is also a boundary case with defs(0): if the returned value is itself a function argument, defs selects the owning function.

Syntax trees, graphs, and sets#

An MLIR document contains several overlapping structures. Zirium keeps the operations in source order, but a query can choose the relationship it needs:

• children and subtree walk nested regions and operations;
• defs and users follow SSA edges;
• parent and root move back through structural ownership;
• closure adds the scopes and symbol definitions needed by a selection;
• reachable follows supported SSA, call, branch, and symbol references.

Those operations compose with unions, intersections, and differences. To find StableHLO matrix multiplications that directly feed an addition, intersect the set of matrix multiplications with the definitions used by additions:

filter(op("stablehlo.dot_general"))
intersect
(filter(op("stablehlo.add")) | defs)

Set operations produce a selection in source order with duplicates removed. Ordinary navigation preserves duplicates, which is useful when counting uses: two paths to the same operation can contribute two items. Use unique when each operation should count only once.

Calls need a little more work because the callee lives elsewhere in the document. Starting from a func.call, a fixed-point closure keeps adding its containing scope, the resolved function, and the function body until nothing changes:

filter(op("func.call")) | fixpoint(closure)

This can extract a dependency fragment from a file containing many unrelated functions. Printing a selection retains its enclosing syntax, but does not guarantee a standalone, valid MLIR program.

Turning a query into a report#

Longer queries can live in .zirium files. Save this as histogram.zirium to count operations in each function body:

functions = filter(op("func.func"));

functions
| map_by(
    attr("sym_name"),
    children | subtree | names | tally
  )
| markdown

For a small example, save this as input.mlir:

module {
  func.func @add_i32(%lhs: i32, %rhs: i32) -> i32 {
    %sum = arith.addi %lhs, %rhs : i32
    func.return %sum : i32
  }

  func.func @square_i32(%value: i32) -> i32 {
    %square = "arith.muli"(%value, %value) : (i32, i32) -> i32
    func.return %square : i32
  }
}

Then run:

zirium -f histogram.zirium input.mlir

The two queries inside map_by run on one function at a time. children | subtree selects its body, including nested operations, without counting the function itself. The result is a GitHub-flavored Markdown table with one row per function and one column per operation name:

| Key | arith.addi | arith.muli | func.return |
| --- | --- | --- | --- |
| add\_i32 | 1 |  | 1 |
| square\_i32 |  | 1 | 1 |

Changing the final stage from markdown to json returns the same counts as a nested object. When several input files are supplied, Zirium evaluates the program independently for each document. Adding --jsonl wraps each emitted result with its input path, which makes the output easier to consume from another program.

The query engine is not tied to the CLI. Rust and Python expose typed builders for the same predicates, navigation steps, and transformations, including both slice directions. They run against the same semantic document and return operation handles, strings, counts, and maps directly to the calling program.

Checking structural expectations#

A query can also check an expectation about the IR and fail if it is not met. For example, a compiler pass may be expected to produce a matrix multiplication and eliminate tensor allocations.

The check assertion operator requires a non-empty result, while check(n) requires an exact number of items:

do filter(op("linalg.matmul"))
   | check("expected a lowered matmul");

do filter(op("bufferization.alloc_tensor"))
   | check(0, "tensor allocations must be eliminated");

Save these checks as checks.zirium and run:

zirium --preset linalg --preset bufferization --silent --strict -f checks.zirium input.mlir

Successful checks produce nothing. A failed check gives a diagnostic and a nonzero exit status. The --strict flag also rejects parsing that leaves incomplete semantic information.

These checks are useful for expectations visible in the IR, such as the elimination of tensor allocations. They cannot establish that a lowering preserves the program’s behaviour.

Editing through the same pipeline#

The CLI currently supports some simple attribute edits. A query can set or remove attributes on its selected operations and then emit the complete document:

zirium \
  'do filter(op("arith.addi")) | set_attr("analysis.tag", "review"); emit' \
  arithmetic.mlir

The do statement suppresses printing of the intermediate selection. The final emit is a separate statement, so it starts with the whole edited document and prints one complete result.

The library offers more control over both edits and output. Rust and Python support transactional semantic edits, including restricted insertion and erasure of operations.

Zirium keeps two representations because inspecting a program and reproducing its source have different requirements. The parser retains the original bytes, tokens, concrete syntax tree, and diagnostics, even for malformed input or invalid UTF-8. The semantic representation makes operations, values, regions, and symbols available for queries and edits.

Library callers can reproduce the original input byte for byte, print canonical MLIR, or use source-preserving output for supported edits. Source-preserving output reuses unchanged text and regenerates changed operations. Inserting or erasing operations discards the source retention needed for that output mode.

The fragment printer tries to preserve the original formatting by default as far as possible. Running zirium 'format' input.mlir normalizes layout while preserving custom assembly, SSA names, aliases, comments, and opaque recovered syntax when available. Use format(assembly = "generic") to write operations in quoted generic MLIR instead.

Looking at compiler changes#

Zirium can also run the query language over a semantic diff:

zirium --diff before.mlir after.mlir

The comparison ignores whitespace, comments, consistent SSA renaming, block labels, and operation locations by default. It reports additions, removals, field changes, rewired operands, and movement. Diff predicates narrow those records:

zirium --diff before.mlir after.mlir \
  'filter(changed("attributes") and dialect("arith")) | json'

The before and after stages project change records back to operations in either input. From there, the usual query stages can inspect users, definitions, parents, or reachable bodies. A query can therefore follow a changed operation to its users without requiring SSA names to match across the two files.

The diff compares represented structure; it does not establish computational equivalence between two programs.

Dialect coverage and limits#

Generic quoted operations are handled directly. Custom assembly forms are harder because each dialect may arrange operands, types, attributes, and regions differently.

Zirium has a declarative registry for selected custom forms, including bundled presets for StableHLO and parts of several LLVM dialects. A project can also load registry descriptions from json. Registration supplies enough structure for parsing, queries, and printing; it does not import ODS or TableGen definitions or promise complete dialect verification.

Unknown custom forms can be retained as syntax with recovery diagnostics. Strict mode rejects incomplete semantic information when a script must not continue with a partial view.

Zirium 0.4.0 targets MLIR 22.1 textual syntax. It does not support MLIR bytecode, full dialect verification, execution, or general IR construction. Jobs that need full dialect semantics still require MLIR itself.

Trying Zirium#

Zirium is just a couple of weeks old and written mostly using GPT 5.6 Sol and some GPT 6 Astra, and it probably has a lot of bugs. However, I have used it for parsing fairly large MLIR files with custom syntax, so it can work.

The Python package is available from PyPI:

python -m pip install zirium

The CLI can be installed from crates.io:

cargo install zirium --locked

Standalone archives are available from GitHub releases. The repository also has worked CLI examples, the complete query language reference, and a capability reference comparing the CLI, Python, and Rust interfaces.