SSA Construction

The SPIR-V accepts programs in Static Single Assigment form. So I wrote SSA -construction into python to SPIR-V compiler spirthon.

Overview

Discovery stage in spirthon generates instruction blocks from pythons bytecode. return, goto, conditional branch are only allowed to end every such block.

SSA requires that every variable in the program is assigned only once. For example in the following program:

i = f()
if i < 10:
    i = 10
return i

It would roughly decompose into following blocks (in pseudo-code notation):

    local i
L0: assign i = f()
    cond i < 10 then L1 else L2
L1: assign i = 10
    goto L2
L2: return i

The local i and assign denotes at non-static assignment. i is assigned twice here, so the above code isn't in SSA form. The variable i must be renamed to attain SSA:

L0: i1 = f()
    cond i1 < 10 then L1 else L2
L1: i2 = 10
    goto L2
L2: i3 = phi(L0: i1, L1: i2)
    return i3

The renaming introduces a Phi -node. In the example it means that goto from L0 causes i1 to be assigned to i3, whereas goto from L1 causes i2 to be assigned to i3.

Basic Blocks

SSA construction may start when the IR is being generated. I avoid generating 'assign' -instructions in spirthon, by holding a table for dependencies and definitions.

Here's an example of code that isn't in SSA-form, yet without assign -instructions:

L0: i1 = f()
    cond i1 < 10 then L1 else L2 
L0.defines = {i: i1}
L0.depends = set()

L1: i2 = 10
    goto L2 
L1.defines = {i: i2}
L1.depends = set()

L2: return i
L2.defines = {}
L2.depends = set([i])

Every basic block here only ever entries and leaves at their ends, so we may substitute away variables with an obvious source, and hold a definition mapping from variables that were defined in that block.

The depends -set is used to recognize variable liveness within the program. We wouldn't want to insert phi -nodes for definitions that are never used anywhere, if we can avoid it.

Variable is live within a block, if it's got a successor that depends on the variable, and if the block doesn't define the variable. In spirthon, we enforce this rule until it is satisfied:

adjust = True
while adjust:
    adjust = False
    for block in proc.blocks:
        L = len(block.depends)
        for succ in block.succ:
            block.depends.update(succ.depends)
            block.depends.difference_update(block.defines)
        if len(block.depends) > L:
            adjust = True

After we know where the variables need to be live, we still do not have enough information to replace remaining variable names with static single assignment -names.

Dominance

Whenever execution flow reaches a block, there exists a set of blocks that must have been visited prior. This set is a dominator set of the block. Immediate dominator is the nearest dominator to the block.

Knowing the dominators in our program graph allows to determine where phi -nodes need to be inserted.

Dominance can be computed from an ordinary depth-first or breath-first search. The immediate dominator of a block lies on the path from entry to the block:

def depth_first_flow(block, prec):
    if block is not entry and block.idom is None:
        block.idom = prec
        for succ in block.succ:
            depth_first_flow(succ, block)
for succ in entry.succ:
    depth_first_flow(succ, entry)

For every block, there is a common immediate dominator that can be found by locating the common block on the paths:

def common_dominator(lhs, rhs):
    lhs_d = distance(lhs)
    rhs_d = distance(rhs)
    while lhs_d > rhs_d:
        lhs = lhs.idom
        lhs_d -= 1
    while rhs_d > lhs_d:
        rhs = rhs.idom
        rhs_d -= 1
    while lhs != rhs:
        lhs = lhs.idom
        rhs = rhs.idom
    return lhs

def distance(block):
    count = 0
    while block.idom is not None:
        block = block.idom
        count += 1
    return count

The immediate dominator of a block is the common immediate dominator of it's preceding blocks. The rule can be enforced:

adjust = True
while adjust:
    adjust = False
    for block in proc.blocks[1:]:
        idom = reduce(common_dominator, block.prec)
        adjust |= (idom <> block.idom)
        block.idom = idom

Here reduce() conveniently expresses exactly what needs to be done.

Dominance Frontiers

Phi -nodes are needed when all of the following conditions match:

Dominance frontier of a block consists of blocks where dominance of the block ends. There is a well-known way to compute dominance frontiers once you know the immediate dominators:

for block in proc.blocks:
    if len(block.prec) >= 2:
        for prec in block.prec:
            runner = prec
            while runner != block.idom:
                runner.frontiers.add(block)
                runner = runner.idom

Inserting Phi -nodes

Once we know the dominance frontiers, we know where to insert phi-nodes too: 

for block in proc.blocks:
    for frontier in block.frontiers:
        for variable in block.defines:
            insert_phi(frontier, variable)

def insert_phi(block, variable):
    if variable in block.depends and variable not in block.phi:
        block.phi[variable] = Phi(block, {})
        for frontier in block.frontiers:
            insert_phi(frontier, variable)

The phi-node insertion produces a new definition, so the code must be written to recurse the phi-insertion for this new definition.

Variable Lookup

Once phi nodes have been filled out, the remaining variables can be substituted away by walking backwards along the immediate dominators.

Note that also the phi-nodes need to be filled out at this point.

for block in proc.blocks:
    for prec in block.prec:
        for variable, phi in block.phi.items():
            phi.args[prec] = lookup(prec, variable)
    for instruction in block.instructions:
        for i, arg in enumerate(instruction.args):
            if isinstance(arg, Local):
                instruction.args[i] = lookup_up(block, arg)

def lookup(block, variable):
    if variable in block.defines:
        value = block.defines[variable]
        if isinstance(value, Local):
            return lookup_up(block, value)
        return value
    return lookup_up(block, variable)

def lookup_up(block, variable):
    if variable in block.phi:
        return block.phi[variable]
    return lookup(block.idom, variable)

Phi-nodes have been inserted on the frontiers, so the relevant phi-nodes end up on the dominator chain of the node.

After there should be no more instructions holding references to local variables.

Fin

This code seems to work for all few situations where I've tried it. But it might still not be correct. Be sure to check the results.

The spirthon is using the exact same code as described above to construct its SSA -form, unless I will find errors and update the code. The latest code may be found in the discovery.py -file inside the spirthon -project.

Similar posts