def fuse_pop_push(program):
"""Fuses pop+push sequences in the given `Program`.
A stack pop followed by a stack push (with no intervening read) is equivalent
to just updating the top of the stack. The latter is more efficient for FULL
variables, because it just updates the cache for the top, and avoids gathering
from and scattering to the backing stack Tensor.
This pass mutates the `ControlFlowGraph` of the input `Program` to convert
pop+push sequences into updates. The pass will work despite intervening
instructions that interact with other Variables, but will not cross basic
block boundaries. As a side-effect, the pass moves non-optimized pops to the
last place in their basic block where they are still sound. This has no
effect on the runtime behavior of the program.
Args:
program: A lowered `Program` whose pop+push sequences to fuse. `Block`s in
the program may be mutated.
Returns:
fused: A `Program` with statically redundant pop+push sequences eliminated
in favor of `PrimOp`s with non-trivial `skip_push_mask` fields.
Raises:
ValueError: If the input `Program` has not been lowered (i.e., contains
`FunctionCallOp`), or is ill-formed (e.g., contains invalid instructions).
"""
for i in range(program.graph.exit_index()):
block = program.graph.block(i)
new_instructions = []
waiting_for_pop = collections.OrderedDict([])
for op in block.instructions:
if isinstance(op, inst.PrimOp):
skip_push = set()
for var in inst.pattern_traverse(op.vars_in):
if var in waiting_for_pop:
new_instructions.append(inst.PopOp([var]))
del waiting_for_pop[var]
for var in inst.pattern_traverse(op.vars_out):
if var in waiting_for_pop:
skip_push.add(var)
del waiting_for_pop[var]
new_instructions.append(
inst.prim_op(op.vars_in, op.vars_out, op.function,
skip_push))
elif isinstance(op, inst.FunctionCallOp):
raise ValueError(
"pop-push fusion not implemented for pre-lowering")
elif isinstance(op, inst.PopOp):
for var in inst.pattern_traverse(op.vars):
if var in waiting_for_pop:
new_instructions.append(inst.PopOp([var]))
else:
waiting_for_pop[var] = True
else:
raise ValueError("Unrecognized op in pop-push fusion", op)
if waiting_for_pop:
new_instructions.append(inst.PopOp(list(waiting_for_pop.keys())))
block.instructions = new_instructions
return program
def _function_entry_stack_ops(op, to_pop):
"""Computes a set of stack operations for the entry to a FunctionCallOp.
The function calling convention is
- Push the arguments to the formal parameters
- Pop any now-dead arguments so they're not on the stack during the call
- Jump to the beginning of the function body
This can be a little tricky for a self-call, because then the arguments and
the formal parameters live in the same name space and can collide. This
helper does something reasonable, and errors out when it can't.
Args:
op: FunctionCallOp instance giving the call to make stack operations for.
to_pop: Set of names to make sure are popped before entering.
Returns:
ops: List of instruction objects that accomplish the goal.
"""
push_from = []
push_to = []
caller_side_vars = inst.pattern_flatten(op.vars_in)
callee_side_vars = inst.pattern_flatten(op.function.vars_in)
for caller_side, callee_side in inst.pattern_zip(caller_side_vars,
callee_side_vars):
if caller_side == callee_side:
# This can happen if this is a self-call and we're just passing the
# variable through to itself
if caller_side in to_pop:
# The top of the stack is already correct, and the callee will pop our
# unneeded value off it for us: skip the push and the pop.
to_pop = to_pop.difference([caller_side])
else:
# The top of the stack is correct, but we need to push it anyway because
# the callee will (eventually) pop but we still need the current value
# when the callee returns.
push_from.append(caller_side)
push_to.append(callee_side)
elif callee_side in caller_side_vars:
# If the graph of transfers turns out not to be a DAG, can't implement it
# without temporary space; and don't want to bother computing a safe order
# even if it does.
# Fixing this is b/135275883.
msg = ('Cannot lower FunctionCallOp that reuses its own input {}'
' as a formal parameter.')
raise ValueError(msg.format(caller_side))
# Checking `elif caller_side in callee_side_vars` is redundant because
# the callee_side check will trigger on that pair sooner or later.
else:
# Ordinary transfer: push now and then pop if needed. The pop will not
# interfere with the push because only caller-side variables can possibly
# be popped.
assert callee_side not in to_pop
push_from.append(caller_side)
push_to.append(callee_side)
push_inst = inst.push_op(push_from, push_to)
if to_pop:
return [push_inst, inst.PopOp(list(to_pop))]
else:
return [push_inst]
def pea_nuts_program(latent_shape, choose_depth, step_state):
"""Synthetic program usable for benchmarking VM performance.
This program is intended to resemble the control flow and scaling
parameters of the NUTS algorithm, without any of the complexity.
Hence the name.
Each batch member looks like:
state = ... # shape latent_shape
def recur(depth, state):
if depth > 1:
state1 = recur(depth - 1, state)
state2 = state1 + 1
state3 = recur(depth - 1, state2)
ans = state3 + 1
else:
ans = step_state(state) # To simulate NUTS, something heavy
return ans
while count > 0:
count = count - 1
depth = choose_depth(count)
state = recur(depth, state)
Args:
latent_shape: Python `tuple` of `int` giving the event shape of the
latent state.
choose_depth: Python `Tensor -> Tensor` callable. The input
`Tensor` will have shape `[batch_size]` (i.e., scalar event
shape), and give the iteration of the outer while loop the
thread is in. The `choose_depth` function must return a `Tensor`
of shape `[batch_size]` giving the depth, for each thread,
to which to call `recur` in this iteration.
step_state: Python `Tensor -> Tensor` callable. The input and
output `Tensor`s will have shape `[batch_size] + latent_shape`.
This function is expected to update the state, and represents
the "real work" versus which the VM overhead is being measured.
Returns:
program: `instructions.Program` that runs the above benchmark.
"""
entry = instructions.Block()
top_body = instructions.Block()
finish_body = instructions.Block()
enter_recur = instructions.Block()
recur_body_1 = instructions.Block()
recur_body_2 = instructions.Block()
recur_body_3 = instructions.Block()
recur_base_case = instructions.Block()
# pylint: disable=bad-whitespace
entry.assign_instructions([
instructions.prim_op(["count"], "cond",
lambda count: count > 0), # cond = count > 0
instructions.BranchOp("cond", top_body,
instructions.halt()), # if cond
])
top_body.assign_instructions([
instructions.PopOp(["cond"]), # done with cond now
instructions.prim_op(["count"], "ctm1",
lambda count: count - 1), # ctm1 = count - 1
instructions.PopOp(["count"]), # done with count now
instructions.push_op(["ctm1"], ["count"]), # count = ctm1
instructions.PopOp(["ctm1"]), # done with ctm1
instructions.prim_op(["count"], "depth",
choose_depth), # depth = choose_depth(count)
instructions.push_op(
["depth", "state"],
["depth", "state"]), # state = recur(depth, state)
instructions.PopOp(["depth", "state"]), # done with depth, state
instructions.PushGotoOp(finish_body, enter_recur),
])
finish_body.assign_instructions([
instructions.push_op(["ans"], ["state"]), # ...
instructions.PopOp(["ans"]), # pop callee's "ans"
instructions.GotoOp(entry), # end of while body
])
# Definition of recur begins here
enter_recur.assign_instructions([
instructions.prim_op(["depth"], "cond1",
lambda depth: depth > 0), # cond1 = depth > 0
instructions.BranchOp("cond1", recur_body_1,
recur_base_case), # if cond1
])
recur_body_1.assign_instructions([
instructions.PopOp(["cond1"]), # done with cond1 now
instructions.prim_op(["depth"], "dm1",
lambda depth: depth - 1), # dm1 = depth - 1
instructions.PopOp(["depth"]), # done with depth
instructions.push_op(
["dm1", "state"],
["depth", "state"]), # state1 = recur(dm1, state)
instructions.PopOp(["state"]), # done with state
instructions.PushGotoOp(recur_body_2, enter_recur),
])
recur_body_2.assign_instructions([
instructions.push_op(["ans"], ["state1"]), # ...
instructions.PopOp(["ans"]), # pop callee's "ans"
instructions.prim_op(["state1"], "state2",
lambda state: state + 1), # state2 = state1 + 1
instructions.PopOp(["state1"]), # done with state1
instructions.push_op(
["dm1", "state2"],
["depth", "state"]), # state3 = recur(dm1, state2)
instructions.PopOp(["dm1", "state2"]), # done with dm1, state2
instructions.PushGotoOp(recur_body_3, enter_recur),
])
recur_body_3.assign_instructions([
instructions.push_op(["ans"], ["state3"]), # ...
instructions.PopOp(["ans"]), # pop callee's "ans"
instructions.prim_op(["state3"], "ans",
lambda state: state + 1), # ans = state3 + 1
instructions.PopOp(["state3"]), # done with state3
instructions.IndirectGotoOp(), # return ans
])
recur_base_case.assign_instructions([
instructions.PopOp(["cond1", "depth"]), # done with cond1, depth
instructions.prim_op(["state"], "ans",
step_state), # ans = step_state(state)
instructions.PopOp(["state"]), # done with state
instructions.IndirectGotoOp(), # return ans
])
pea_nuts_graph = instructions.ControlFlowGraph([
entry,
top_body,
finish_body,
enter_recur,
recur_body_1,
recur_body_2,
recur_body_3,
recur_base_case,
])
# pylint: disable=bad-whitespace
pea_nuts_vars = {
"count": instructions.single_type(np.int64, ()),
"cond": instructions.single_type(np.bool, ()),
"cond1": instructions.single_type(np.bool, ()),
"ctm1": instructions.single_type(np.int64, ()),
"depth": instructions.single_type(np.int64, ()),
"dm1": instructions.single_type(np.int64, ()),
"state": instructions.single_type(np.float32, latent_shape),
"state1": instructions.single_type(np.float32, latent_shape),
"state2": instructions.single_type(np.float32, latent_shape),
"state3": instructions.single_type(np.float32, latent_shape),
"ans": instructions.single_type(np.float32, latent_shape),
}
return instructions.Program(pea_nuts_graph, [], pea_nuts_vars,
["count", "state"], "state")
def is_even_function_calls(include_types=True, dtype=np.int64):
"""The is-even program, via "even-odd" recursion.
Computes True if the input is even, False if the input is odd, by a pair of
mutually recursive functions is_even and is_odd, which return True and False
respectively for <1-valued inputs.
Tests out mutual recursion.
Args:
include_types: If False, we omit types on the variables, requiring a type
inference pass.
dtype: The dtype to use for `n`-like internal state variables.
Returns:
program: Full-powered `instructions.Program` that computes is_even(n).
"""
def pred_type(t):
return instructions.TensorType(np.bool, t[0].shape)
# Forward declaration of is_odd.
is_odd_func = instructions.Function(None, ["n"], "ans", pred_type)
enter_is_even = instructions.Block()
finish_is_even = instructions.Block()
recur_is_even = instructions.Block()
is_even_func = instructions.Function(None, ["n"], "ans", pred_type)
# pylint: disable=bad-whitespace
# Definition of is_even function
enter_is_even.assign_instructions([
instructions.prim_op(["n"], "cond", lambda n: n < 1), # cond = n < 1
instructions.BranchOp("cond", finish_is_even,
recur_is_even), # if cond
])
finish_is_even.assign_instructions([
instructions.PopOp(["n", "cond"]), # done with n, cond
instructions.prim_op([], "ans", lambda: True), # ans = True
instructions.halt_op(), # return ans
])
recur_is_even.assign_instructions([ # else
instructions.PopOp(["cond"]), # done with cond now
instructions.prim_op(["n"], "nm1", lambda n: n - 1), # nm1 = n - 1
instructions.PopOp(["n"]), # done with n
instructions.FunctionCallOp(is_odd_func, ["nm1"],
"ans"), # ans = is_odd(nm1)
instructions.PopOp(["nm1"]), # done with nm1
instructions.halt_op(), # return ans
])
is_even_blocks = [enter_is_even, finish_is_even, recur_is_even]
is_even_func.graph = instructions.ControlFlowGraph(is_even_blocks)
enter_is_odd = instructions.Block()
finish_is_odd = instructions.Block()
recur_is_odd = instructions.Block()
# pylint: disable=bad-whitespace
# Definition of is_odd function
enter_is_odd.assign_instructions([
instructions.prim_op(["n"], "cond", lambda n: n < 1), # cond = n < 1
instructions.BranchOp("cond", finish_is_odd, recur_is_odd), # if cond
])
finish_is_odd.assign_instructions([
instructions.PopOp(["n", "cond"]), # done with n, cond
instructions.prim_op([], "ans", lambda: False), # ans = False
instructions.halt_op(), # return ans
])
recur_is_odd.assign_instructions([ # else
instructions.PopOp(["cond"]), # done with cond now
instructions.prim_op(["n"], "nm1", lambda n: n - 1), # nm1 = n - 1
instructions.PopOp(["n"]), # done with n
instructions.FunctionCallOp(is_even_func, ["nm1"],
"ans"), # ans = is_even(nm1)
instructions.PopOp(["nm1"]), # done with nm1
instructions.halt_op(), # return ans
])
is_odd_blocks = [enter_is_odd, finish_is_odd, recur_is_odd]
is_odd_func.graph = instructions.ControlFlowGraph(is_odd_blocks)
is_even_main_blocks = [
instructions.Block([
instructions.FunctionCallOp(is_even_func, ["n1"], "ans"),
], instructions.halt_op()),
]
# pylint: disable=bad-whitespace
is_even_vars = {
"n": instructions.single_type(dtype, ()),
"n1": instructions.single_type(dtype, ()),
"cond": instructions.single_type(np.bool, ()),
"nm1": instructions.single_type(dtype, ()),
"ans": instructions.single_type(np.bool, ()),
}
if not include_types:
_strip_types(is_even_vars)
return instructions.Program(
instructions.ControlFlowGraph(is_even_main_blocks),
[is_even_func, is_odd_func], is_even_vars, ["n1"], "ans")
def fibonacci_program():
"""More complicated, fibonacci program: computes fib(n): fib(0) = fib(1) = 1.
Returns:
program: Full-powered `instructions.Program` that computes fib(n).
"""
entry = instructions.Block(name="entry")
enter_fib = instructions.Block(name="enter_fib")
recur1 = instructions.Block(name="recur1")
recur2 = instructions.Block(name="recur2")
recur3 = instructions.Block(name="recur3")
finish = instructions.Block(name="finish")
# pylint: disable=bad-whitespace
entry.assign_instructions([
instructions.PushGotoOp(instructions.halt(), enter_fib),
])
# Definition of fibonacci function starts here
enter_fib.assign_instructions([
instructions.prim_op(["n"], "cond", lambda n: n > 1), # cond = n > 1
instructions.BranchOp("cond", recur1, finish), # if cond
])
recur1.assign_instructions([
instructions.PopOp(["cond"]), # done with cond now
instructions.prim_op(["n"], "nm1", lambda n: n - 1), # nm1 = n - 1
instructions.push_op(["nm1"], ["n"]), # fibm1 = fibonacci(nm1)
instructions.PopOp(["nm1"]), # done with nm1
instructions.PushGotoOp(recur2, enter_fib),
])
recur2.assign_instructions([
instructions.push_op(["ans"], ["fibm1"]), # ...
instructions.PopOp(["ans"]), # pop callee's "ans"
instructions.prim_op(["n"], "nm2", lambda n: n - 2), # nm2 = n - 2
instructions.PopOp(["n"]), # done with n
instructions.push_op(["nm2"], ["n"]), # fibm2 = fibonacci(nm2)
instructions.PopOp(["nm2"]), # done with nm2
instructions.PushGotoOp(recur3, enter_fib),
])
recur3.assign_instructions([
instructions.push_op(["ans"], ["fibm2"]), # ...
instructions.PopOp(["ans"]), # pop callee's "ans"
instructions.prim_op(["fibm1", "fibm2"], "ans",
lambda x, y: x + y), # ans = fibm1 + fibm2
instructions.PopOp(["fibm1", "fibm2"]), # done with fibm1, fibm2
instructions.IndirectGotoOp(), # return ans
])
finish.assign_instructions([ # else:
instructions.PopOp(["n", "cond"]), # done with n, cond
instructions.prim_op([], "ans", lambda: 1), # ans = 1
instructions.IndirectGotoOp(), # return ans
])
fibonacci_blocks = [entry, enter_fib, recur1, recur2, recur3, finish]
# pylint: disable=bad-whitespace
fibonacci_vars = {
"n": instructions.single_type(np.int64, ()),
"cond": instructions.single_type(np.bool, ()),
"nm1": instructions.single_type(np.int64, ()),
"fibm1": instructions.single_type(np.int64, ()),
"nm2": instructions.single_type(np.int64, ()),
"fibm2": instructions.single_type(np.int64, ()),
"ans": instructions.single_type(np.int64, ()),
}
return instructions.Program(
instructions.ControlFlowGraph(fibonacci_blocks), [], fibonacci_vars,
["n"], "ans")
def _lower_function_calls_1(builder,
graph,
defined_in,
live_out,
function=True):
"""Lowers one function body, destructively.
Mutates the given `ControlFlowGraphBuilder`, inserting `Block`s
representing the new body. Some of these may be the same as some
`Block`s in the input `graph`, mutated; others may be newly
allocated.
Args:
builder: `ControlFlowGraphBuilder` constructing the answer.
graph: The `ControlFlowGraph` to lower.
defined_in: A Python list of `str`. The set of variables that
are defined on entry to this `graph`.
live_out: A Python list of `str`. The set of variables that are
live on exit from this `graph`.
function: Python `bool`. If `True` (the default), assume this is
a `Function` body and convert an "exit" transfer into
`IndirectGotoOp`; otherwise leave it as (`Program`) "exit".
Raises:
ValueError: If an invalid instruction is encountered, if a live
variable is undefined, if different paths into a `Block` cause
different sets of variables to be defined, or if trying to lower
function calls in a program that already has `IndirectGotoOp`
instructions (they confuse the liveness analysis).
"""
liveness_map = liveness.liveness_analysis(graph, set(live_out))
defined_map = _definedness_analysis(graph, defined_in, liveness_map)
for i in range(graph.exit_index()):
block = graph.block(i)
old_instructions = block.instructions
# Resetting block.instructions here because we will build up the
# list of new ones in place (via the `builder`).
block.instructions = []
builder.append_block(block)
builder.maybe_add_pop(defined_map[block].defined_into_block,
liveness_map[block].live_into_block)
for op_i, (op, defined_out, live_out) in enumerate(
zip(old_instructions,
defined_map[block].defined_out_instructions,
liveness_map[block].live_out_instructions)):
if isinstance(op, inst.PrimOp):
for name in inst.pattern_traverse(op.vars_in):
if name in inst.pattern_flatten(op.vars_out):
# Why not? Because the stack discipline we are trying to
# implement calls for popping variables as soon as they
# become dead. Now, if a PrimOp writes to the same
# variable as it reads, the old version of that variable
# dies. Where to put the PopOp? Before the PrimOp is no
# good -- it still needs to be read. After the PrimOp is
# no good either -- it will pop the output, not the input.
# Various solutions to this problem are possible, such as
# adding a "drop the second-top element of this stack"
# instruction, or orchestrating the pushes and pops
# directly in the interpreter, but for now the simplest
# thing is to just forbid this situation.
# Fixing this is b/118884528.
msg = 'Cannot lower PrimOp that writes to its own input {}.'
raise ValueError(msg.format(name))
builder.append_instruction(op)
builder.maybe_add_pop(defined_out, live_out)
elif isinstance(op, inst.FunctionCallOp):
names_pushed_here = inst.pattern_flatten(op.vars_out)
for name in inst.pattern_traverse(op.vars_in):
if name in names_pushed_here:
# For the same reason as above.
# Fixing this is b/118884528.
msg = 'Cannot lower FunctionCallOp that writes to its own input {}.'
raise ValueError(msg.format(name))
# The variables that were defined on entry to this instruction (i.e.,
# not pushed here) but are not live out don't need to remain on their
# stacks when the callee is entered.
defined_in = defined_out.difference(names_pushed_here)
to_pop = defined_in.difference(live_out)
for new_op in _function_entry_stack_ops(op, to_pop):
builder.append_instruction(new_op)
if (op_i == len(old_instructions) - 1
and _optimizable_tail_call(op, builder.cur_block())):
builder.end_block_with_tail_call(
op.function.graph.block(0))
# The check that the tail call is optimizable is equivalent to
# checking that the push-pop pair below would do nothing.
else:
builder.split_block(op.function.graph.block(0))
builder.append_instruction(
# These extra levels of list protect me (I hope) from the
# auto-unpacking in the implementation of push_op, in the case of
# a function returning exactly one Tensor.
inst.push_op([op.function.vars_out], [op.vars_out]))
builder.append_instruction(
inst.PopOp(inst.pattern_flatten(op.function.vars_out)))
# The only way this would actually add a pop is if some name written
# by this call was a dummy variable.
builder.maybe_add_pop(frozenset(names_pushed_here),
live_out)
elif isinstance(op, (inst.PopOp)):
# Presumably, lowering is applied before any `PopOp`s are present. That
# said, simply propagating them is sound. (But see the `PopOp` case in
# `liveness_analysis`.)
builder.append_instruction(op)
else:
raise ValueError('Invalid instruction in block {}.'.format(op))
if function:
builder.maybe_adjust_terminator()
def maybe_add_pop(self, defined, live):
poppable = defined.difference(live)
if poppable:
self.append_instruction(inst.PopOp(list(poppable)))