modelverse/kernel/modelverse_jit/cfg_optimization.py

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

from collections import defaultdict
import modelverse_jit.cfg_ir as cfg_ir
import modelverse_jit.cfg_dominators as cfg_dominators

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

        # Branches to blocks which contain nothing but a jump can be replaced by branches
        # to the jump's target.
        for branch in block.flow.branches():
            if (is_empty_block(branch.block)
                    and branch.block is not block
                    and isinstance(branch.block.flow, cfg_ir.JumpFlow)):
                new_branch = branch.block.flow.branch
                branch.block = new_branch.block
                branch.arguments = new_branch.arguments
                changed = True

def get_all_blocks(entry_point):
    """Gets all basic blocks in the control-flow graph defined by the given entry point."""
    yield entry_point
    for block in get_reachable_blocks(entry_point):
        yield block

def optimize_graph_flow(entry_point):
    """Optimizes all flow instructions in the graph defined by the given entry point."""
    for block in get_all_blocks(entry_point):
        optimize_flow(block)

def merge_blocks(entry_point):
    """Merges blocks which have exactly one predecessor with said predecessor, if the
       predecessor has a jump flow instruction."""
    predecessor_map = cfg_dominators.get_all_predecessor_blocks(entry_point)
    queue = list(predecessor_map.keys())
    def __do_merge(source, target):
        for target_param, branch_arg in zip(target.parameters, source.flow.branch.arguments):
            source.append_definition(target_param)
            target_param.redefine(branch_arg)

        for target_def in target.definitions:
            source.append_definition(target_def)

        source.flow = target.flow
        for pred_set in predecessor_map.values():
            if target in pred_set:
                pred_set.remove(target)
                pred_set.add(source)

    while len(queue) > 0:
        block = queue.pop()
        preds = predecessor_map[block]
        if len(preds) == 1:
            single_pred = next(iter(preds))
            if isinstance(single_pred.flow, cfg_ir.JumpFlow):
                __do_merge(single_pred, block)

def elide_local_checks(entry_point):
    """Tries to elide redundant checks on local variables."""
    # The plan here is to replace all check-local-exists defs by literals if
    # they are either dominated by an appropriate declare-local or not reachable
    # from a declare-local.
    local_checks = defaultdict(set)
    local_defs = defaultdict(set)
    for block in get_all_blocks(entry_point):
        for definition in block.definitions:
            if cfg_ir.is_value_def(definition, cfg_ir.CheckLocalExists):
                local_checks[cfg_ir.get_def_variable(definition).node_id].add(definition)
            elif cfg_ir.is_value_def(definition, cfg_ir.DeclareLocal):
                local_defs[cfg_ir.get_def_variable(definition).node_id].add(definition)

    dominator_tree = cfg_dominators.get_dominator_tree(entry_point)
    reachable_blocks = get_all_reachable_blocks(entry_point)
    for (variable, all_checks) in local_checks.items():
        for check in all_checks:
            is_reachable = False
            for local_def in local_defs[variable]:
                if dominator_tree.dominates_instruction(local_def, check):
                    # Check is dominated by a definition. Replace it by a 'True' literal.
                    check.redefine(cfg_ir.Literal(True))
                    is_reachable = True
                    break
                elif check.block in reachable_blocks[local_def.block]:
                    is_reachable = True

            if not is_reachable:
                # Check cannot be reached from any definition. Replace it by a 'False' literal.
                check.redefine(cfg_ir.Literal(False))

def eliminate_unused_definitions(entry_point):
    """Tries to eliminate unused definitions in the control-flow graphb defined by the
       given entry point."""
    def_dependencies = {}
    root_defs = set()
    for block in get_all_blocks(entry_point):
        for definition in block.definitions:
            def_dependencies[definition] = set(
                [dep for dep in definition.get_all_dependencies()
                 if isinstance(dep, cfg_ir.Definition)])
            if definition.has_side_effects():
                root_defs.add(definition)

        for dep in block.flow.get_all_dependencies():
            if isinstance(dep, cfg_ir.Definition):
                root_defs.add(dep)

    live_defs = set()
    def __mark_live(definition):
        if definition in live_defs:
            return

        live_defs.add(definition)
        if definition in def_dependencies:
            for dep in def_dependencies[definition]:
                __mark_live(dep)

    for root in root_defs:
        __mark_live(root)

    dead_defs = set.difference(set(def_dependencies.keys()), live_defs)
    for dead_def in dead_defs:
        dead_def.block.remove_definition(dead_def)

def try_redefine_as_direct_call(definition, jit, called_globals):
    """Tries to redefine the given indirect call definition as a direct call."""
    call = definition.value
    if not isinstance(call, cfg_ir.IndirectFunctionCall):
        return

    target = cfg_ir.get_def_value(call.target)
    if isinstance(target, cfg_ir.LoadPointer):
        loaded_ptr = cfg_ir.get_def_value(target.pointer)
        if isinstance(loaded_ptr, cfg_ir.ResolveGlobal):
            resolved_var_name = loaded_ptr.variable.name

            # # Try to resolve the callee as an intrinsic.
            # intrinsic = jit.get_intrinsic(resolved_var_name)
            # if intrinsic is not None:
            #     return redefine_as_intrinsic(definition, intrinsic, call.argument_list)

            # Otherwise, build a thunk.
            thunk_name = jit.jit_thunk_global(resolved_var_name)
            definition.redefine(
                cfg_ir.DirectFunctionCall(
                    thunk_name, call.argument_list, cfg_ir.JIT_CALLING_CONVENTION))
            called_globals.add(loaded_ptr)
    elif isinstance(target, cfg_ir.Literal):
        node_id = target.literal
        thunk_name = jit.jit_thunk_constant(node_id)
        definition.redefine(
            cfg_ir.DirectFunctionCall(
                thunk_name, call.argument_list, cfg_ir.JIT_CALLING_CONVENTION))

def optimize_calls(entry_point, jit):
    """Converts indirect calls to direct calls in the control-flow graph defined by the
       given entry point."""
    called_globals = set()
    for block in get_all_blocks(entry_point):
        for definition in block.definitions:
            try_redefine_as_direct_call(definition, jit, called_globals)

def optimize(entry_point, jit):
    """Optimizes the control-flow graph defined by the given entry point."""
    optimize_graph_flow(entry_point)
    elide_local_checks(entry_point)
    optimize_graph_flow(entry_point)
    optimize_calls(entry_point, jit)
    eliminate_unused_definitions(entry_point)
    optimize_graph_flow(entry_point)
    merge_blocks(entry_point)