295 lines
9.8 KiB
Python
295 lines
9.8 KiB
Python
|
"""Transformation for inserting refrecence count inc/dec opcodes.
|
||
|
|
|
||
|
|
This transformation happens towards the end of compilation. Before this
|
||
|
|
transformation, reference count management is not explicitly handled at all.
|
||
|
|
By postponing this pass, the previous passes are simpler as they don't have
|
||
|
|
to update reference count opcodes.
|
||
|
|
|
||
|
|
The approach is to decrement reference counts soon after a value is no
|
||
|
|
longer live, to quickly free memory (and call __del__ methods), though
|
||
|
|
there are no strict guarantees -- other than that local variables are
|
||
|
|
freed before return from a function.
|
||
|
|
|
||
|
|
Function arguments are a little special. They are initially considered
|
||
|
|
'borrowed' from the caller and their reference counts don't need to be
|
||
|
|
decremented before returning. An assignment to a borrowed value turns it
|
||
|
|
into a regular, owned reference that needs to freed before return.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from typing import Dict, Iterable, Tuple
|
||
|
|
|
||
|
|
from mypyc.analysis.dataflow import (
|
||
|
|
AnalysisDict,
|
||
|
|
analyze_borrowed_arguments,
|
||
|
|
analyze_live_regs,
|
||
|
|
analyze_must_defined_regs,
|
||
|
|
cleanup_cfg,
|
||
|
|
get_cfg,
|
||
|
|
)
|
||
|
|
from mypyc.ir.func_ir import FuncIR, all_values
|
||
|
|
from mypyc.ir.ops import (
|
||
|
|
Assign,
|
||
|
|
BasicBlock,
|
||
|
|
Branch,
|
||
|
|
ControlOp,
|
||
|
|
DecRef,
|
||
|
|
Goto,
|
||
|
|
IncRef,
|
||
|
|
Integer,
|
||
|
|
KeepAlive,
|
||
|
|
LoadAddress,
|
||
|
|
Op,
|
||
|
|
Register,
|
||
|
|
RegisterOp,
|
||
|
|
Value,
|
||
|
|
)
|
||
|
|
|
||
|
|
Decs = Tuple[Tuple[Value, bool], ...]
|
||
|
|
Incs = Tuple[Value, ...]
|
||
|
|
|
||
|
|
# A cache of basic blocks that decrement and increment specific values
|
||
|
|
# and then jump to some target block. This lets us cut down on how
|
||
|
|
# much code we generate in some circumstances.
|
||
|
|
BlockCache = Dict[Tuple[BasicBlock, Decs, Incs], BasicBlock]
|
||
|
|
|
||
|
|
|
||
|
|
def insert_ref_count_opcodes(ir: FuncIR) -> None:
|
||
|
|
"""Insert reference count inc/dec opcodes to a function.
|
||
|
|
|
||
|
|
This is the entry point to this module.
|
||
|
|
"""
|
||
|
|
cfg = get_cfg(ir.blocks)
|
||
|
|
values = all_values(ir.arg_regs, ir.blocks)
|
||
|
|
|
||
|
|
borrowed = {value for value in values if value.is_borrowed}
|
||
|
|
args: set[Value] = set(ir.arg_regs)
|
||
|
|
live = analyze_live_regs(ir.blocks, cfg)
|
||
|
|
borrow = analyze_borrowed_arguments(ir.blocks, cfg, borrowed)
|
||
|
|
defined = analyze_must_defined_regs(ir.blocks, cfg, args, values, strict_errors=True)
|
||
|
|
ordering = make_value_ordering(ir)
|
||
|
|
cache: BlockCache = {}
|
||
|
|
for block in ir.blocks.copy():
|
||
|
|
if isinstance(block.ops[-1], (Branch, Goto)):
|
||
|
|
insert_branch_inc_and_decrefs(
|
||
|
|
block,
|
||
|
|
cache,
|
||
|
|
ir.blocks,
|
||
|
|
live.before,
|
||
|
|
borrow.before,
|
||
|
|
borrow.after,
|
||
|
|
defined.after,
|
||
|
|
ordering,
|
||
|
|
)
|
||
|
|
transform_block(block, live.before, live.after, borrow.before, defined.after)
|
||
|
|
|
||
|
|
cleanup_cfg(ir.blocks)
|
||
|
|
|
||
|
|
|
||
|
|
def is_maybe_undefined(post_must_defined: set[Value], src: Value) -> bool:
|
||
|
|
return isinstance(src, Register) and src not in post_must_defined
|
||
|
|
|
||
|
|
|
||
|
|
def maybe_append_dec_ref(
|
||
|
|
ops: list[Op], dest: Value, defined: AnalysisDict[Value], key: tuple[BasicBlock, int]
|
||
|
|
) -> None:
|
||
|
|
if dest.type.is_refcounted and not isinstance(dest, Integer):
|
||
|
|
ops.append(DecRef(dest, is_xdec=is_maybe_undefined(defined[key], dest)))
|
||
|
|
|
||
|
|
|
||
|
|
def maybe_append_inc_ref(ops: list[Op], dest: Value) -> None:
|
||
|
|
if dest.type.is_refcounted:
|
||
|
|
ops.append(IncRef(dest))
|
||
|
|
|
||
|
|
|
||
|
|
def transform_block(
|
||
|
|
block: BasicBlock,
|
||
|
|
pre_live: AnalysisDict[Value],
|
||
|
|
post_live: AnalysisDict[Value],
|
||
|
|
pre_borrow: AnalysisDict[Value],
|
||
|
|
post_must_defined: AnalysisDict[Value],
|
||
|
|
) -> None:
|
||
|
|
old_ops = block.ops
|
||
|
|
ops: list[Op] = []
|
||
|
|
for i, op in enumerate(old_ops):
|
||
|
|
key = (block, i)
|
||
|
|
|
||
|
|
assert op not in pre_live[key]
|
||
|
|
dest = op.dest if isinstance(op, Assign) else op
|
||
|
|
stolen = op.stolen()
|
||
|
|
|
||
|
|
# Incref any references that are being stolen that stay live, were borrowed,
|
||
|
|
# or are stolen more than once by this operation.
|
||
|
|
for j, src in enumerate(stolen):
|
||
|
|
if src in post_live[key] or src in pre_borrow[key] or src in stolen[:j]:
|
||
|
|
maybe_append_inc_ref(ops, src)
|
||
|
|
# For assignments to registers that were already live,
|
||
|
|
# decref the old value.
|
||
|
|
if dest not in pre_borrow[key] and dest in pre_live[key]:
|
||
|
|
assert isinstance(op, Assign)
|
||
|
|
maybe_append_dec_ref(ops, dest, post_must_defined, key)
|
||
|
|
|
||
|
|
# Strip KeepAlive. Its only purpose is to help with this transform.
|
||
|
|
if not isinstance(op, KeepAlive):
|
||
|
|
ops.append(op)
|
||
|
|
|
||
|
|
# Control ops don't have any space to insert ops after them, so
|
||
|
|
# their inc/decrefs get inserted by insert_branch_inc_and_decrefs.
|
||
|
|
if isinstance(op, ControlOp):
|
||
|
|
continue
|
||
|
|
|
||
|
|
for src in op.unique_sources():
|
||
|
|
# Decrement source that won't be live afterwards.
|
||
|
|
if src not in post_live[key] and src not in pre_borrow[key] and src not in stolen:
|
||
|
|
maybe_append_dec_ref(ops, src, post_must_defined, key)
|
||
|
|
# Decrement the destination if it is dead after the op and
|
||
|
|
# wasn't a borrowed RegisterOp
|
||
|
|
if (
|
||
|
|
not dest.is_void
|
||
|
|
and dest not in post_live[key]
|
||
|
|
and not (isinstance(op, RegisterOp) and dest.is_borrowed)
|
||
|
|
):
|
||
|
|
maybe_append_dec_ref(ops, dest, post_must_defined, key)
|
||
|
|
block.ops = ops
|
||
|
|
|
||
|
|
|
||
|
|
def insert_branch_inc_and_decrefs(
|
||
|
|
block: BasicBlock,
|
||
|
|
cache: BlockCache,
|
||
|
|
blocks: list[BasicBlock],
|
||
|
|
pre_live: AnalysisDict[Value],
|
||
|
|
pre_borrow: AnalysisDict[Value],
|
||
|
|
post_borrow: AnalysisDict[Value],
|
||
|
|
post_must_defined: AnalysisDict[Value],
|
||
|
|
ordering: dict[Value, int],
|
||
|
|
) -> None:
|
||
|
|
"""Insert inc_refs and/or dec_refs after a branch/goto.
|
||
|
|
|
||
|
|
Add dec_refs for registers that become dead after a branch.
|
||
|
|
Add inc_refs for registers that become unborrowed after a branch or goto.
|
||
|
|
|
||
|
|
Branches are special as the true and false targets may have a different
|
||
|
|
live and borrowed register sets. Add new blocks before the true/false target
|
||
|
|
blocks that tweak reference counts.
|
||
|
|
|
||
|
|
Example where we need to add an inc_ref:
|
||
|
|
|
||
|
|
def f(a: int) -> None
|
||
|
|
if a:
|
||
|
|
a = 1
|
||
|
|
return a # a is borrowed if condition is false and unborrowed if true
|
||
|
|
"""
|
||
|
|
prev_key = (block, len(block.ops) - 1)
|
||
|
|
source_live_regs = pre_live[prev_key]
|
||
|
|
source_borrowed = post_borrow[prev_key]
|
||
|
|
source_defined = post_must_defined[prev_key]
|
||
|
|
|
||
|
|
term = block.terminator
|
||
|
|
for i, target in enumerate(term.targets()):
|
||
|
|
# HAX: After we've checked against an error value the value we must not touch the
|
||
|
|
# refcount since it will be a null pointer. The correct way to do this would be
|
||
|
|
# to perform data flow analysis on whether a value can be null (or is always
|
||
|
|
# null).
|
||
|
|
omitted: Iterable[Value]
|
||
|
|
if isinstance(term, Branch) and term.op == Branch.IS_ERROR and i == 0:
|
||
|
|
omitted = (term.value,)
|
||
|
|
else:
|
||
|
|
omitted = ()
|
||
|
|
|
||
|
|
decs = after_branch_decrefs(
|
||
|
|
target, pre_live, source_defined, source_borrowed, source_live_regs, ordering, omitted
|
||
|
|
)
|
||
|
|
incs = after_branch_increfs(target, pre_live, pre_borrow, source_borrowed, ordering)
|
||
|
|
term.set_target(i, add_block(decs, incs, cache, blocks, target))
|
||
|
|
|
||
|
|
|
||
|
|
def after_branch_decrefs(
|
||
|
|
label: BasicBlock,
|
||
|
|
pre_live: AnalysisDict[Value],
|
||
|
|
source_defined: set[Value],
|
||
|
|
source_borrowed: set[Value],
|
||
|
|
source_live_regs: set[Value],
|
||
|
|
ordering: dict[Value, int],
|
||
|
|
omitted: Iterable[Value],
|
||
|
|
) -> tuple[tuple[Value, bool], ...]:
|
||
|
|
target_pre_live = pre_live[label, 0]
|
||
|
|
decref = source_live_regs - target_pre_live - source_borrowed
|
||
|
|
if decref:
|
||
|
|
return tuple(
|
||
|
|
(reg, is_maybe_undefined(source_defined, reg))
|
||
|
|
for reg in sorted(decref, key=lambda r: ordering[r])
|
||
|
|
if reg.type.is_refcounted and reg not in omitted
|
||
|
|
)
|
||
|
|
return ()
|
||
|
|
|
||
|
|
|
||
|
|
def after_branch_increfs(
|
||
|
|
label: BasicBlock,
|
||
|
|
pre_live: AnalysisDict[Value],
|
||
|
|
pre_borrow: AnalysisDict[Value],
|
||
|
|
source_borrowed: set[Value],
|
||
|
|
ordering: dict[Value, int],
|
||
|
|
) -> tuple[Value, ...]:
|
||
|
|
target_pre_live = pre_live[label, 0]
|
||
|
|
target_borrowed = pre_borrow[label, 0]
|
||
|
|
incref = (source_borrowed - target_borrowed) & target_pre_live
|
||
|
|
if incref:
|
||
|
|
return tuple(
|
||
|
|
reg for reg in sorted(incref, key=lambda r: ordering[r]) if reg.type.is_refcounted
|
||
|
|
)
|
||
|
|
return ()
|
||
|
|
|
||
|
|
|
||
|
|
def add_block(
|
||
|
|
decs: Decs, incs: Incs, cache: BlockCache, blocks: list[BasicBlock], label: BasicBlock
|
||
|
|
) -> BasicBlock:
|
||
|
|
if not decs and not incs:
|
||
|
|
return label
|
||
|
|
|
||
|
|
# TODO: be able to share *partial* results
|
||
|
|
if (label, decs, incs) in cache:
|
||
|
|
return cache[label, decs, incs]
|
||
|
|
|
||
|
|
block = BasicBlock()
|
||
|
|
blocks.append(block)
|
||
|
|
block.ops.extend(DecRef(reg, is_xdec=xdec) for reg, xdec in decs)
|
||
|
|
block.ops.extend(IncRef(reg) for reg in incs)
|
||
|
|
block.ops.append(Goto(label))
|
||
|
|
cache[label, decs, incs] = block
|
||
|
|
return block
|
||
|
|
|
||
|
|
|
||
|
|
def make_value_ordering(ir: FuncIR) -> dict[Value, int]:
|
||
|
|
"""Create a ordering of values that allows them to be sorted.
|
||
|
|
|
||
|
|
This omits registers that are only ever read.
|
||
|
|
"""
|
||
|
|
# TODO: Never initialized values??
|
||
|
|
result: dict[Value, int] = {}
|
||
|
|
n = 0
|
||
|
|
|
||
|
|
for arg in ir.arg_regs:
|
||
|
|
result[arg] = n
|
||
|
|
n += 1
|
||
|
|
|
||
|
|
for block in ir.blocks:
|
||
|
|
for op in block.ops:
|
||
|
|
if (
|
||
|
|
isinstance(op, LoadAddress)
|
||
|
|
and isinstance(op.src, Register)
|
||
|
|
and op.src not in result
|
||
|
|
):
|
||
|
|
# Taking the address of a register allows initialization.
|
||
|
|
result[op.src] = n
|
||
|
|
n += 1
|
||
|
|
if isinstance(op, Assign):
|
||
|
|
if op.dest not in result:
|
||
|
|
result[op.dest] = n
|
||
|
|
n += 1
|
||
|
|
elif op not in result:
|
||
|
|
result[op] = n
|
||
|
|
n += 1
|
||
|
|
|
||
|
|
return result
|