|View source on GitHub|
An auto-batchable program.
The primary ingredient of a Program is the control flow graph of operations to perform. The operation language is a union that serves two purposes: one subset is designed to be convenient to run in Single Instruction Multiple Thread style, and the other to generate from an upstream Python-embedded DSL.
As such, there are operations for explicit control transfers and stack management, as well as for interpreted function calls (pending lowering to explicit control transfers). The primitive computations are encapsulated from the interpreter as Python functions. It is not expected that one would author programs in this operation language directly, but rather generate them with an appropriate compiler.
Lowering consists of eliminating
FunctionCallOp in favor of a
specific sequence of lower-level instructions. A few choices for
the lowered language may register as somewhat nonstandard:
- The variable name space is global; the upstream compiler is
expected to generate unique variable names.
- There is no one stack; instead, every variable has its own stack.
This reduces push traffic, because only the variables that are
actually written to are ever pushed.
- Return addresses for pending lowered function calls are stored in
a reserved variable, that is otherwise in the same environment as
the user variables.
- There are no designated registers for function arguments or return
values. This is because all runtime storage is in Tensors, which
need to have fixed types. Instead, it is the (post-lowering)
caller's responsibility to write the arguments into the formal
parameters and to retrieve the returned value(s) from the
variable(s) in which the callee left them.
The post-lowering function call sequence is
- Push the arguments to the formal parameters;
- Pop any argument variables that are no longer used;
- Store the desired return address and jump to the beginning of the function's
body (with a single
- When the function returns by executing
IndirectGotoOp, assign the
returned values to the variables that should receive them; and
- Pop the variables holding the returned values.
Note that this sequence requires that all calls in the source language be to statically known functions, and for every function to leave its results in the same variable(s) on every call (regardless of internal control flow).
__init__( graph, functions, var_defs, vars_in, vars_out, var_alloc=None )
Initialize a new
ControlFlowGraph. This is the graph of basic blocks to execute.
functions: A list of
Functions giving the definitions of all the auto-batchable functions this
Programmay (recursively) call.
var_defs: A dict mapping variable names to
Typeobjects giving their pattern of held Tensors. Each leaf of the pattern is a
TensorTypeobject giving the dtype and shape of that leaf. The shape excludes the batch dimension.
vars_in: A list of the names of the variables in which to store the inputs when starting.
vars_out: A pattern of the names of the variables from which to read the outputs when finished.
var_alloc: A dict mapping variable names to allocation strategies (see
VariableAllocation). The semantics of an entry are "A proof has been found that this strategy suffices for this variable."
Return a representation of the main program as a
replace( var_defs=None, var_alloc=None )
Return a copy of