How to iterate all block statements in IRCode?

Internals & Design
1

Say, I have an instance of IRCode like this:

7 1 ─ %1 = Main.bar(_2)::Tuple{Float64, Float64}                                      │
  │   %2 = Base.indexed_iterate(%1, 1)::PartialStruct(Tuple{Float64, Int64}, Any[Float64, Const(2)])
  │   %3 = Core.getfield(%2, 1)::Float64                                              │
  │   %4 = Core.getfield(%2, 2)::Const(2)                                             │
  │   %9 = π (%4, Const(2))                                                           │
  │   %5 = Base.indexed_iterate(%1, 2, %9)::PartialStruct(Tuple{Float64, Int64}, Any[Float64, Const(3)])
  │   %6 = Core.getfield(%5, 1)::Float64                                              │
8 │   %7 = (%3 + %6)::Float64                                                         │
  └──      return %7

I know how to:

  • get the list of all original nodes: ir.stmts.inst
  • get the ids of nodes in the ith block: ir.cfg.blocks[i].stmts
  • get the list of new nodes, e.g. the PiNode above, as well as their place in the common list of nodes: ir.new_nodes.stmts.inst, ir.new_nodes.info

Now I want to get the list of all nodes - the original and the new ones - belonging to a particular block. Is there a mapping of blocks to new nodes or an utility function to iterate all block statements in the correct order?

2

I might be wrong but I believe IncrementalCompact does something like what you described

3

Indeed, the code looks relevant, but I can’t figure out how to use it to get all nodes for a particular block. Simply iterating over the IncrementalCompact also doesn’t include new nodes:

julia> for node in Core.Compiler.IncrementalCompact(ir) println(node) end
Pair{Pair{Int64, Int64}, Any}(1 => 1, :(Main.bar(_2)))
Pair{Pair{Int64, Int64}, Any}(2 => 2, :(Base.indexed_iterate(%1, 1)))
Pair{Pair{Int64, Int64}, Any}(3 => 3, :(Core.getfield(%2, 1)))
Pair{Pair{Int64, Int64}, Any}(4 => 4, :(Core.getfield(%2, 2)))
Pair{Pair{Int64, Int64}, Any}(5 => 5, :(Base.indexed_iterate(%1, 2, %4)))
Pair{Pair{Int64, Int64}, Any}(6 => 6, :(Core.getfield(%5, 1)))
Pair{Pair{Int64, Int64}, Any}(7 => 7, :(%3 + %6))
Pair{Pair{Int64, Int64}, Any}(8 => 8, :(return %7))

Is there a chance you have a relevant example?

4

Check out Core.Compiler.bbidxiter and CFG (which is stored in the IRCode, I believe)

5

bbidxiter seems to be even closer, but I’m still missing the exact way to utilize it:

julia> itr = Core.Compiler.bbidxiter(ir)
7 1 ─ %1 = Main.bar(_2)::Tuple{Float64, Float64}                                      │
  │   %2 = Base.indexed_iterate(%1, 1)::PartialStruct(Tuple{Float64, Int64}, Any[Float64, Const(2)])
  │   %3 = Core.getfield(%2, 1)::Float64                                              │
  │   %4 = Core.getfield(%2, 2)::Const(2)                                             │
  │        π (%4, Const(2))      │   # <-- PiNode in the correct position!
  │   %5 = Base.indexed_iterate(%1, 2, %4)::PartialStruct(Tuple{Float64, Int64}, Any[Float64, Const(3)])
  │   %6 = Core.getfield(%5, 1)::Float64                                              │
8 │   %7 = (%3 + %6)::Float64                                                         │
  └──      return %7                                                                  │
  ) 


julia> Core.Compiler.iterate(itr, (1, 1))
((1, 1), (2, 1))

julia> Core.Compiler.iterate(itr)
((1, 1), (2, 1))

julia> Core.Compiler.iterate(itr, (2, 1))
((1, 2), (3, 1))

julia> Core.Compiler.iterate(itr, (3, 1))
((1, 3), (4, 1))

julia> Core.Compiler.iterate(itr, (4, 1))
((1, 4), (5, 1))

# I'd expect the PiNode to arise around this place
julia> Core.Compiler.iterate(itr, (5, 1))
((1, 5), (6, 1))

julia> Core.Compiler.iterate(itr, (6, 1))
((1, 6), (7, 1))

julia> Core.Compiler.iterate(itr, (7, 1))
((1, 7), (8, 1))

julia> Core.Compiler.iterate(itr, (8, 1))
((1, 8), (9, 2))

julia> Core.Compiler.iterate(itr, (9, 1))

On the other hand, the answers in this thread made me look into the implementation of show_ir() which prints the nodes in correct order. It seems like they iterate over original nodes and pop new nodes when needed, regardless of the current block. I’m going to try the same strategy and see how it works for my case.

6

For bbidxiter, I’d just toss it in a for loop, check when the block number (first element) matches the one(s) you and and then index into the IRCode with statement number (second element). A little easier but maybe less robust would be to index BasicBlocks on IRCode.CFG directly and iterate over their stmts.

Note that I’m not sure how well/if at all these work on non-compacted IR. There’s a bit of handling of range wraparound on the basic block side, but it’s likely far from bulletproof.