ZJIT adds support for Iongraph, which offers a web-based, pass-by-pass viewer with a stable layout, better navigation, and quality-of-life features like labeled backedges and clickable operands.

Prelude

I’m an intern on the ZJIT team for the fall term. I also have a rather bad habit of being chronically on lobste.rs.

While idly browsing, I spotted an article by Ben Visness titled Who needs Graphviz when you can build it yourself?, which covers his work on creating a novel graph viewer called Iongraph.

Iongraph used to visualize an IR graph for the SpiderMonkey JIT inside Firefox.

Immediately, I was intrigued. I like looking at new technology and I wondered what it might be like to integrate the work done in Iongraph with ZJIT, getting all sorts of novel and interesting features for free. I suspected that it could also help other engineers to reason about how their optimizations might affect the control flow graph of a given function.

Also, it just looks really cool. It’s got nice colours, good built-in CSS, and is built in a fairly extensible way. The underlying code isn’t hard to read if you need to make changes to it.

Investigating further

Iongraph is compelling for a few reasons.

It supports stable layouts, which means that removing or adding nodes (something that can happen when you run an optimization pass) doesn’t shift the location of other nodes to an extreme degree. Iongraph also gives all sorts of interactive options, like clickable operands, scrollable graphs, or arrow keys to navigate between different nodes.

An especially useful feature is the ability to switch between different compiled methods with a small selector. In our codebase, ZJIT compiles each method on its own, so using a tool like this allows us to inspect method level optimizations all in one pane of a web browser. Of course, there are other great features, like loop header highlighting or being able to click on optimization passes to see what the control flow graph looks like after they’re applied.

Proposal

Roughly an hour after I read through said article, I noticed that my mentor, Max, had also posted it in an internal team chat, mentioning that it would be cool to support it.

Of course, I was tempted by this project. As is a common trait for interns, I was tempted to take on a new, shiny project despite not knowing what it might imply to actually develop it. After talking to Max further, he clarified that this would require significant infrastructure work — or at the very least, more than initially apparent.

Building

A JSON library inside ZJIT?

Looking into the Iongraph format, I figured that I would have to use some sort of JSON crate. Since ZJIT as a project doesn’t rely strictly on using Rust tooling like cargo, directly adding serde_json as a dependency was out of the question. Another compelling option was vendoring it (or a smaller JSON library), but that was likely to include features that we did not want or introduce licensing issues.

After a quick discussion, I settled on implementing the functionality myself. I read a bit of the JSON specification, and got a sense of the ideal way to design the library’s API. Ultimately, I chose to opt for readability and usability over raw performance. This decision I think is reasonable given that the serialization code is not in the critical path of the compiler. It’s also accurate to say that the interface is clean enough to replace the internals in the future with minimal issue should there be more performance needed.

In designing the serializer, I chose to target RFC 8259, which provides more freedom than previous specifications. As noted in said RFC, historical specifications constrained the top level value to be an array or an object, but this spec (and my implementation) don’t require that constraint. I also opted to avoid comments, encode strictly in UTF-8, and escape control characters. Notably, RFC 8259 does not impose a limit on precision of numbers, just that infinity, negative infinity, or NaN are restricted.

Computing control flow graph properties

With JSON serialization handled, the more challenging work was computing the graph metadata that Iongraph requires. The format expects explicit successor and predecessor relationships, loop headers, and back edge sources — information that ZJIT doesn’t normally compute since it’s not needed for compilation at this stage of compiler development.

One constraint I had to contend with was that the Iongraph format needs the user to manually provide the successor and predecessor nodes for a given node in a control flow graph. In ZJIT, we compile individual methods at a time as Functions (our internal representation) that hold a graph of Blocks. Each Block is a basic block that you would find in a compiler textbook. (One caveat to understand is that we use extended basic blocks, meaning that blocks can have jump instructions at any point in their contained instructions — not just at the end.)

The process of computing successors and predecessors is fairly simple. As you iterate through the list of blocks, all blocks referenced as the target of a jump-like instruction (whether conditional or unconditional) are added to the successor set. Then for each successor, update their predecessor set to include the block currently being operated on.

The next task I had to solve was computing the loop headers and back edge sources.

Required in the process of computing both of these are computing the dominators for blocks in a control flow graph. We can state that a block i dominates a block j if all paths in the control flow graph that reach j must go through i. Several algorithms exist for computing dominators. There exist both simple iterative options and more complicated versions. Initially, I heard of a fixed point iteration option that was very straightforward to implement but perhaps not the most efficient. That one (which I will discuss shortly) runs in quadratic time to the number of blocks available. In A Simple, Fast Dominance Algorithm by Cooper, Harvey, and Kennedy, both this iterative solution as well as one that is optimized to use less space are mentioned. A third option is the Lengauer-Tarjan algorithm, which has better worst case bounds compared to both the iterative and tuned implementations.

Based on the goals of the project, I opted to use the iterative algorithm, since it performs well and doesn’t incur serious memory use penalties for a small number of blocks in a control flow graph. It can be described as such:

dom = {}
nodes.each do |node|
  if entry_nodes.include?(node)
    dom[node] = Set[node]
  else
    dom[node] = nodes.to_set
  end
end

changed = true
while changed
  changed = false
  nodes.reverse_post_order.each do |node|
    preds = predecessors(node)
    pred_doms = preds.map { |p| dom[p] }

    # Intersection of all predecessor dominators
    intersection = if pred_doms.empty?
                     Set.new
                   else
                     pred_doms.reduce(:&)
                   end

    # Union with {node}
    new_set = intersection | Set[node]

    if new_set != dom[node]
      dom[node] = new_set
      changed = true
    end
  end
end

Implementing this is fairly simple, and it runs quickly enough for the limited number of nodes that it is totally acceptable.

To compute successors we use the following snippet:

let successors: BTreeSet<BlockId> = block
    .insns
    .iter()
    .map(|&insn_id| uf.find_const(insn_id))
    .filter_map(|insn_id| {
        Self::extract_jump_target(&function.insns[insn_id.0])
    })
    .collect();

Here we go through all the instructions in a given block. We use a union find data structure to map instructions to their canonical representatives (since some optimizations may have merged or aliased instructions). We then filter by extract_jump_target, which returns an Option that contains a BlockId for jump-like instructions.

After finding successors, we can set the predecessors by iterating through the nodes in the successor set and adding the current node to their predecessor sets.

The last important thing we need to consider is finding the loop depth.

For finding this, we need to consider first how we can find a natural loop in the first place.

We identify natural loops by detecting back edges. A back edge occurs when a block has a predecessor that is dominated by that block (all paths to the predecessor pass through this block). When we find such an edge, the target block is a loop header and the predecessor is the source of a back edge. The natural loop consists of all blocks on paths from the back edge source to the loop header (excluding the header itself). Each block within this natural loop then has its loop depth incremented.

These additional computations are used within the Iongraph layout engine to determine the height at which a given block should be vertically, or where lines should be routed within the graph. Loop headers and back edge sources are also marked!

The final result

You can click around this demo graph showing a simple example from ZJIT to get a sense of how Iongraph works! Operands are clickable to get to their definition. You can click on the phases of optimization on the left side - note that only the non-grayed out passes will have made changes. The graph is also zoomable and scrollable!

Hopefully this post was educational! I learned a lot implementing this feature and enjoyed doing so.

If you would like to do some work on ZJIT (and learn a lot in the process), you are welcome to make pull requests to github.com/ruby/ruby/ with the commit prefix ZJIT:. You can find issues here.

Also, feel free to join our Zulip!