modelverse/kernel/modelverse_jit/cfg_optimization.py

"""Optimizes and analyzes CFG-IR."""

import modelverse_jit.cfg_ir as cfg_ir

def get_directly_reachable_blocks(block):
    """Gets the set of all blocks that can be reached by taking a single branch from the
       given block."""
    return [branch.block for branch in block.flow.branches()]

def get_reachable_blocks(entry_point):
    """Constructs the set of all reachable vertices from the given block."""
    # This is a simple O(n^2) algorithm. Maybe a faster algorithm is more appropriate here.
    def __add_block_children(block, results):
        for child in get_directly_reachable_blocks(block):
            if child not in results:
                results.add(child)
                __add_block_children(child, results)

        return results

    return __add_block_children(entry_point, set())

def get_all_reachable_blocks(entry_point):
    """Constructs the set of all reachable vertices, for every block that is
       reachable from the given entry point."""
    # This is a simple O(n^3) algorithm. Maybe a faster algorithm is more appropriate here.
    results = {}
    all_blocks = get_reachable_blocks(entry_point)
    results[entry_point] = all_blocks
    for block in all_blocks:
        if block not in results:
            results[block] = get_reachable_blocks(block)
    return results

def is_empty_block(block):
    """Tests if the given block contains no parameters or definitions."""
    return len(block.parameters) == 0 and len(block.definitions) == 0

def optimize_flow(block):
    """Optimizes the given block's flow instruction."""
    changed = True
    while changed:
        changed = False

        # Select flow with a literal condition can be optimized to a direct jump.
        if (isinstance(block.flow, cfg_ir.SelectFlow)
                and cfg_ir.is_literal_def(block.flow.condition)):
            literal = cfg_ir.get_literal_def_value(block.flow.condition)
            block.flow = cfg_ir.JumpFlow(
                block.flow.if_branch if literal else block.flow.else_branch)
            changed = True

        # Jumps to blocks which contain no parameters or definitions can be replaced
        # by the target block's flow.
        if (isinstance(block.flow, cfg_ir.JumpFlow)
                and is_empty_block(block.flow.branch.block)
                and block.flow.branch.block is not block):
            block.flow = block.flow.branch.block.flow
            changed = True

def optimize_graph_flow(entry_point):
    """Optimizes all flow instructions in the graph defined by the given entry point."""
    all_blocks = set(get_reachable_blocks(entry_point))
    all_blocks.add(entry_point)
    for block in all_blocks:
        optimize_flow(block)

def optimize(entry_point):
    """Optimizes the control-flow graph defined by the given entry point."""
    optimize_graph_flow(entry_point)