Querying and editing MLIR from the command line

Pass a query followed by input paths. Each file is an independent MLIR document; omitting paths reads stdin. From a source checkout, build the CLI with:

cargo build --bin zirium

The commands below assume target/debug is on PATH:

export PATH="$PWD/target/debug:$PATH"

The query language reference lists every predicate, pipeline stage, and output rule.

Inputs and reusable .zirium programs are in examples/cli. Short queries are shown inline.

The default baseline registry accepts ordinary, named, and nested module shorthand. Use --preset stablehlo for a bundled dialect or --registry FILE for a custom configuration. Both flags are repeatable. See custom formats for details.

zirium --help lists options, and zirium --list-presets lists available dialects. Recovery warnings on stderr mean semantic queries may be incomplete. By default, reachable retains an operation with unknown reference semantics as a leaf and continues elsewhere. Add --strict to reject either case in scripts. Options can appear before or after the query argument. Use -- before input paths beginning with a dash.

Troubleshoot source errors

Syntax and lowering failures are written to stderr with the input or query name, a one-based line and display column, a specific message, a bounded source excerpt, and a marker. For example, an unresolved operand looks like:

$ zirium count broken.mlir
zirium: could not lower broken.mlir:
broken.mlir:2:17: error: unresolved SSA value `%missing`
  "example.use"(%missing) : (i32) -> ()
                ^~~~~~~~

Inline queries use <query> as their source name; -f diagnostics use the program path. Columns count Unicode characters rather than UTF-8 bytes. Invalid UTF-8 bytes are shown as �, each invalid byte occupies one column, tabs advance to four-column stops, CR in CRLF is omitted, and an error at end of file gets a marker at the corresponding empty or final line. Source lines longer than 120 display columns are clipped around the marker and marked with …. The original byte ranges remain attached to parser and lowering diagnostics. As with other failures before delivery, stdout stays empty.

Every CLI evaluation provides {document} as the supplied input path, or stdin for standard input. Use it in report text or JSON strings:

zirium 'print("# {document}"); names | tally | markdown' first.mlir second.mlir

For machine-readable multi-file output, --jsonl writes one compact record per emitted result with document and result fields. JSON results remain JSON values; text and selected MLIR are strings. Records keep input and statement order, and stdout remains empty if any input or emission fails. --ndjson is an alias for --jsonl:

zirium --jsonl 'names | tally | json' first.mlir second.mlir

Find untagged arithmetic

Boolean predicates can combine operation names and attributes. This command prints the untagged arith.addi from the sample and omits its tagged arith.muli:

zirium \
  'filter((op("arith.addi") or op("arith.muli")) and not has_attr("analysis.tag"))' \
  examples/cli/arithmetic.mlir

The result contains the add and its enclosing module syntax. Its operand definitions and users are omitted, so the fragment may not be valid on its own. The same query is available as untagged-arithmetic.zirium.

Find direct consumers

users follows direct SSA uses. This sample’s add has one consumer, the multiply:

$ cat examples/cli/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
}
$ zirium 'filter(op("arith.addi")) | users' examples/cli/arithmetic.mlir
builtin.module {
  %v3 = "arith.muli"(%v2, %v1) {analysis.tag = "old"} : (i32, i32) -> i32
}

The same query is available as direct-consumers.zirium.

Combine arithmetic operation kinds

union combines the results of two queries. The result contains the add and multiply in source order:

zirium \
  'filter(op("arith.addi")) union filter(op("arith.muli"))' \
  examples/cli/arithmetic.mlir

This is also a selected fragment, so its SSA dependencies are omitted. The same query is available as arithmetic-union.zirium.

Extract a call dependency slice

fixpoint(closure) adds dependencies until the selection stops changing. For a func.call, that includes the resolved callee and its body, and leaves unrelated sibling functions out.

zirium 'filter(op("func.call")) | fixpoint(closure)' examples/cli/calls.mlir

The output contains @caller and @answer, and omits @unrelated. The slice may still need other parts of the input to be valid MLIR. The same query is available as call-closure.zirium.

Count operations in a StableHLO decoder

The larger stablelm-decode.mlir example represents one token-generation step for a small two-layer StableLM. It contains rotary position handling, KV-cache updates, grouped-query attention, gated MLPs, and a language-model head.

The example is adapted from MLXcel’s Apache-2.0 licensed StableLM decode program. Its large embedded rotary tables are replaced by zero splats; the operation structure and tensor shapes are unchanged.

The first query below counts every operation. The second counts reductions:

$ zirium --preset stablehlo --strict 'count' examples/cli/stablelm-decode.mlir
237
$ zirium --preset stablehlo --strict 'filter(op("stablehlo.reduce")) | count' examples/cli/stablelm-decode.mlir
14

The initial selection contains every operation, including the module and other container operations.

For a query that will be reused, put the program in a file. The checked-in stablehlo-matmul-count.zirium counts stablehlo.dot_general, the StableHLO operation used for the decoder’s matrix multiplications:

$ cat examples/cli/stablehlo-matmul-count.zirium
filter(op("stablehlo.dot_general")) | count
$ zirium --preset stablehlo --strict --program-file examples/cli/stablehlo-matmul-count.zirium examples/cli/stablelm-decode.mlir
19

The nineteen operations comprise nine matrix multiplications in each decoder layer and one final projection to logits. A program file contains only Zirium source; surrounding whitespace and the final newline are ignored.

The count is structural: compact reductions using applies stablehlo.add do not synthesize reducer-body operations. Use an explicit-region input if those operations must participate in queries.

Inspect shapes and operation kinds

Name and type projections make inventories easy to pipe into ordinary tools:

zirium --preset stablehlo --strict \
  'filter(dialect("stablehlo")) | names' examples/cli/stablelm-decode.mlir \
  | sort | uniq -c | sort -nr

zirium --preset stablehlo --strict \
  'filter(op("stablehlo.dot_general")) | result_types | unique' \
  examples/cli/stablelm-decode.mlir

Use result_type("tensor<32xf32>") to select that exact result type, or json to inspect operand types, result types, and complete dimension clauses together. Paired custom dimensions such as [0] x [1] remain intact in attribute projection.

Inspect a single output computation

The decoder returns logits, a key cache, and a value cache. Select the first return operand and follow its SSA definitions, stopping at function inputs:

zirium --preset stablehlo --strict \
  'filter(op("func.return")) | defs(0) | slice' \
  examples/cli/stablelm-decode.mlir

Change the index to 1 or 2 for either cache. slice does not expand symbols or region bodies. At a multi-result operation it follows all explicit operands, without inferring result-specific dependencies. The output remains an inspection fragment. Use fixpoint(closure) when retaining scopes and supported callees is more important than a narrow slice.

Keep operations satisfying a relationship

Find matmuls that directly feed an addition by intersecting the matmul set with the definitions used by additions:

zirium --preset stablehlo --strict \
  '(filter(op("stablehlo.dot_general")) intersect (filter(op("stablehlo.add")) | defs)) | json' \
  examples/cli/stablelm-decode.mlir

This retains ten matmuls in the decoder. Use users(0) to follow only the first result of a multi-result operation.

Tag selected operations

Edits preserve the current selection. A do statement suppresses the edited selection’s implicit output; the following emit starts from the whole edited document and prints all operations once:

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

The output contains the complete document, with analysis.tag = "review" added to arith.addi. The same query is available as tag-add.zirium.

Remove a tag while reading standard input

When no input path follows the query, the binary reads MLIR from standard input:

zirium \
  'do filter(string_attr_eq("analysis.tag", "old")) | remove_attr("analysis.tag"); emit' \
  < examples/cli/arithmetic.mlir

The output retains the complete document but removes analysis.tag from arith.muli. The same query is available as remove-tag.zirium.

Inspect an intermediate selection

emit prints the selection and passes it to the next stage. This prints the add, then its direct users, back-to-back:

zirium 'filter(op("arith.addi")) | emit | users' examples/cli/arithmetic.mlir

Zirium does not insert separators between statement or intermediate outputs. Use print("text") when a report needs an explicit heading or delimiter.

Selection output includes enclosing operation shells. Use --fragment-scope minimal when large incidental attributes on those ancestors would dominate a small fragment. The selected operations and selected bodies keep their attributes; enclosing shells retain structural identity and signature metadata. Fragments may still lack SSA definitions, callees, or terminators needed for valid standalone MLIR.

Work inside a function fragment

subtree includes every descendant of the selected function. The following filter therefore finds returns only inside @caller:

zirium \
  'filter(op("func.call")) | parent | subtree | filter(op("func.return"))' \
  examples/cli/calls.mlir

Per-function operation inventories

Run the saved histogram query on the decoder:

zirium --preset stablehlo --strict \
  -f examples/cli/function-op-counts.zirium examples/cli/stablelm-decode.mlir

The result is a JSON object keyed by function name, containing counts keyed by operation name. In this example, main contains 19 stablehlo.dot_general operations. Add reachable | before names to include supported callees’ bodies, counting shared operations once. See bindings and aggregation for the evaluation rules and the distinction between body and reachable counts.

For a Markdown report with an interpolated function count and explanatory text:

zirium --preset stablehlo --strict \
  -f examples/cli/function-report.zirium examples/cli/stablelm-decode.mlir

See Markdown and text reports for supported table shapes and interpolation rules.

To package the histograms with metadata in one JSON document:

zirium --preset stablehlo --strict \
  -f examples/cli/function-report-json.zirium examples/cli/stablelm-decode.mlir

The literal syntax preserves numeric counts and nested maps while allowing interpolated report labels.

Compare two documents

zirium --diff examples/cli/diff-before.mlir examples/cli/diff-after.mlir
zirium --diff examples/cli/diff-before.mlir examples/cli/diff-after.mlir \
  -f examples/cli/diff-report.zirium

The first command prints a human summary. The second emits the changed arith attribute records as zirium.diff.v1 JSON.