clika_runtime.graph
The graph as data.
A compiled or traced :class:ModelGraph answers nodes() with its
:class:Node views: each carries an :class:OpCode, a name, its attributes
as plain Python values, and the :class:Value views of the tensors it
consumes and produces; bound_tensors() names the weights a traced module
bound into the operator (constants() lists edge constants such as an ONNX
initializer). find_nodes(op_code), node(name), constants() and
is_finalized() complete the read surface, and a
:class:Pattern (a graph of operator constraints) finds its occurrences
through find_pattern as :class:Match entries::
graph = crt.io.OnnxModel.open("model.onnx").compile()
for node in graph.nodes():
print(node.name, node.op_code, node.attributes)
(relu,) = graph.find_nodes(crt.graph.OpCode.Relu)
for match in graph.find_pattern(crt.graph.Pattern([OpCode.MatMul, OpCode.Relu])):
matmul, relu = match
Holding a view keeps its graph alive. A view carries its key (a node's
name, a value's producer and port, an edge's two nodes and ports), and after
an edit such as optimize() it finds its node, value or edge again by that
key. Once the key is gone (an edit removed or renamed it, or a merge, a
split or an extract consumed the graph), check() and every other
accessor raise :class:~clika_runtime.InvalidArgumentError naming it, and
so does a query given the view; ==, hash and repr keep working,
so a dict keyed by views still finds its entries.
The graph is a directed multigraph over its nodes. An :class:Edge is one
read of a value, from the producer's output port to the reader's input
port, so Mul(x, x) reads x over two edges; a constant operand is a
:class:Value but no edge. edges() lists every edge and
edges_between(src, dst) the edges from one node to another. A node
answers in_edges(), out_edges(), input(port), output(port),
in_degree() and out_degree(), and a value its uses() and
consumers(). find_nodes also takes a predicate over a :class:Node,
and find_nodes_by_name takes a regular expression matched against the
whole name or, with NameMatch.Anywhere, some part of it.
to_networkx() returns the graph as a networkx.MultiDiGraph for the
algorithms networkx provides: its nodes are the :class:Node views, which
keep this graph alive, and each edge is keyed (src_port, dst_port)::
sinks = graph.find_nodes(lambda node: node.out_degree() == 0)
for edge in graph.edges():
print(edge.src().name, edge.src_port(), edge.dst().name, edge.dst_port())
nx_graph = graph.to_networkx() # needs the networkx package
A walk goes breadth-first from a node: walk(start) lists the nodes it
reaches, ancestors(node) and descendants(node) the nodes before and
after one, and visit(start, visitor) shows a callable each node, which
answers a :class:VisitAction to go on (Continue or None), not go past
the node (Prune) or end the walk (Stop). Keyword arguments bound and
filter every walk: direction (a :class:WalkDirection), max_depth,
max_nodes (a budget the walk refuses to exceed rather than stopping
short), stop_at and edge_filter (callables over a :class:Node and an
:class:Edge) and skip_shape_lane. topological_order lists every node
after the nodes it reads, in the :class:TopologicalOrder asked for;
random_topological_order(seed) draws one of the valid orders;
layers() (or topological_generations()) and depths() give each
node's generation; and is_topological_order checks an order::
relu = graph.node("relu")
upstream = graph.ancestors(relu, max_depth=2)
region = graph.walk(relu, stop_at=lambda node: node.op_code == crt.graph.OpCode.Add)
order = graph.topological_order(crt.graph.TopologicalOrder.MinMemory)
assert graph.is_topological_order(order)
A path never repeats a node. all_simple_paths(src, dst) and
all_simple_edge_paths(src, dst) list the paths between two nodes (as
nodes, where parallel edges give one path, or as edges, where each gives its
own), find_paths_to(start, target) the paths that end at the first node a
callable accepts, and count_paths counts them without listing them. The
listings search at the call and return an iterator; without a limit, a
search that finds more than MAX_PATHS_WITHOUT_LIMIT paths raises
:class:~clika_runtime.InvalidArgumentError rather than returning a cut list.
has_path and shortest_path answer reachability and the fewest edges.
Keyword arguments bound every path query: cutoff (the most edges),
avoid (nodes no path may hold) and edge_filter. Given an
edge_cost, shortest_path returns the cheapest path as a
:class:WeightedPath carrying its nodes, edges and cost, and
critical_path(node_cost, edge_cost) returns the costliest one::
paths = list(graph.all_simple_paths(first, last, cutoff=6))
cheapest = graph.shortest_path(first, last, edge_cost=lambda edge: 1.0 + edge.dst_port())
longest = graph.critical_path(node_cost=lambda node: len(node.outputs))
print(cheapest.cost, [node.name for node in longest.nodes])
The structure of the graph follows from its paths. dominators() maps each
node to the last node every path from the inputs to it crosses, and
post_dominators() to the first node every path from it to the outputs
crosses. split_points() lists the nodes every input-to-output path
crosses, where the graph splits into two stages. region(entry, exit) is
the single-entry, single-exit block between two nodes, between(src, dst)
every node on a path from one to the other, producer_cone and
consumer_cone everything a set of nodes reads or feeds, and
is_convex tells whether a set can be cut out as one piece.
lowest_common_ancestors and merge_point find where paths come from
and where they meet; weakly_connected_components, find_cycle and
would_create_cycle answer connectivity. min_value_cut(before, after)
finds the split of the nodes into two parts that hands the fewest bytes
across, with before and after pinning nodes to either part, and
returns a :class:ValueCut of the values that cross and their bytes::
first, second = graph.split_points()[:2]
block = graph.region(first, second) # the nodes between two split points
assert graph.is_convex(block)
cut = graph.min_value_cut()
print(cut.bytes, [value.name for value in cut.values])
The graph also answers what shapes it admits and how fast it runs:
shape_domains() reports one :class:ShapeDomain per input (each
dimension :class:DimDomain Fixed, Range or Free), validate_shapes
checks shapes against them, and bench times synthesized or given inputs
on a finalized graph into a :class:BenchReport::
domains = graph.shape_domains()
graph.finalize()
report = graph.bench([[8, 131]], crt.graph.BenchOptions(warmup=2, iterations=10))
print(report.p50_ms, report.throughput_items_per_s)
print(report.to_csv())
Before finalize(), the graph can be edited in place.
replace_all_uses_with, set_input, set_constant_input and
bypass_node rewire the reads of a value, add_constant makes a constant
to wire in, add_node makes an operator from its :class:OpCode, its input
values and its settings (under the names Node.attributes reports),
remove_node removes an operator nothing reads, and rename_node
renames one; add_input, add_output, remove_output,
rename_input and rename_output edit the graph's inputs and outputs.
insert(fn, inputs) adds the operators a Python function computes from
values of the graph, traced over one tensor per value, and
replace(old_outputs, fn, inputs) puts them in the place of the operators
that computed old_outputs; an exception fn raises leaves the edit as
that same exception.
Each edit checks everything before it changes anything, so a refused edit (an
:class:~clika_runtime.InvalidArgumentError naming the call and what it
refused, with the code_name FAILED_PRECONDITION on a finalized graph)
leaves the graph exactly as it was. A value takes another's place only with
the same dtype and dims, and the reads of a state buffer, the weights an
operator holds itself and the names a KV cache layer binds stay as they are::
(relu,) = graph.find_nodes(crt.graph.OpCode.Relu)
graph.bypass_node(relu) # the Relu's readers read its input instead
graph.finalize()
Several graphs become one with merge(parts, connections). The parts are
ModelGraphs under names, a dict in its order or a sequence of (name, graph) pairs, and each connection feeds an output of one part into an
input of another, both ends written "<part>.<name>". Parts no connection
joins sit side by side. Every operator of part p is named p/<name>;
io_names says how the inputs and outputs left over are named
("prefix_on_collision", "prefix_always" or "keep"), and
cross_device whether a connection between two devices moves its value
("move") or refuses ("refuse"). The merge returns the merged graph
and one :class:Rename per name it changed, whose kind is a
:class:NameKind. ModelGraph.merge(second, io_map) merges second
into a graph in place and returns the renames. Both consume what they merge
in. On success each part, and second, is left empty and its views raise;
a refused merge leaves every graph as it was::
merged, renames = crt.graph.merge({"enc": encoder, "dec": decoder}, [("enc.h", "dec.x")])
for rename in renames:
print(rename.part, rename.kind, rename.from_, rename.to)
merged.finalize()
A graph comes apart with split. split(graph, cut) cuts it in two at
a :class:ValueCut, which min_value_cut() returns, or at a list of
values: the part "before" holds the operators that produce the cut
values and every operator they depend on, and "after" the rest.
split(graph, partition) names each operator's part with a function over
its :class:Node, called once per operator in nodes() order; inside it
the graph answers reads and queries, while an edit, a run, optimize()
or a merge of it is refused. Both return the parts, a dict of name to
ModelGraph, and the connections that join them, the shapes merge takes,
so merge(*split(graph, ...)) gives back a graph with the source's inputs
and outputs. extract(graph, inputs, outputs) takes the operators that
compute outputs from inputs as a graph of its own and frees the
weights of the operators it leaves out. Each consumes the graph it is
given, as a merge does; a refusal leaves the graph as it was, and an
exception the partition function raises comes back as itself::
parts, connections = crt.graph.split(source, source.min_value_cut())
rebuilt, _ = crt.graph.merge(parts, connections) # the source's inputs and outputs again
joined, _ = crt.graph.merge({"enc": encoder, "dec": decoder}, [("enc.h", "dec.x")])
parts, connections = crt.graph.split(joined, lambda node: node.name.split("/")[0]) # enc, dec
relu = graph.node("relu")
block = crt.graph.extract(graph, [relu.input(0)], [relu.output(0)])
compile() and trace() return the graph as built, and a graph runs
once it is finalized. optimize() runs the graph optimizer first when
wanted, takes its choices as a :class:DefaultsOptions, and returns an
:class:OptimizeReport::
graph = crt.io.OnnxModel.open("model.onnx").compile()
report = graph.optimize(max_num_iterations=10)
graph.finalize()
The optimizer's transforms are named by :class:Transform handles, and a
Python function that edits the graph can be one too (Transform(name, fn)
or the transforms.transform decorator), which optimize() runs beside
the runtime's own. So can a rule: transforms.rewrite(name, pattern, replacement, condition) (or the transforms.rule decorator) replaces
each occurrence of a :class:Pattern its condition accepts with the values
replacement computes, handing it a :class:RewriteMatch and one traced
tensor per value the occurrence reads.
crt.graph.transforms has one factory per built-in transform, named as the
transform is, and transforms.defaults(options) returns the default list
for a set of choices::
handle = crt.graph.transforms.fusion_driver()
print(handle.name) # fusion_driver
optimize(transforms) runs a list of them in order each iteration, and
default_transforms(defaults) returns the list it runs when handed none, a
starting point to edit. A list that places a transform ahead of one it must
follow, or holds one the choices switch off for this graph, is refused before
anything runs. The report carries one :class:TransformReport per list
entry: how often the transform changed the graph, whether it failed, and
whether it was still firing when max_num_iterations stopped the run::
graph = crt.io.OnnxModel.open("model.onnx").compile()
listed = [entry for entry in graph.default_transforms() if entry.name != "fusion_driver"]
report = graph.optimize(listed)
for row in report.transforms:
print(row.transform.name, row.applications, row.failed)
| Name | Kind |
|---|---|
BenchOptions | class |
BenchReport | class |
DefaultsOptions | class |
DimDomain | class |
Edge | class |
FeedFill | class |
InputShape | class |
Match | class |
ModelGraph | class |
NameKind | class |
NameMatch | class |
Node | class |
NodeKind | class |
OpCode | class |
OptimizeReport | class |
Pattern | class |
Rename | class |
RewriteMatch | class |
ShapeDomain | class |
TopologicalOrder | class |
Transform | class |
TransformKind | class |
TransformReport | class |
Value | class |
ValueCut | class |
VisitAction | class |
WalkDirection | class |
WeightedPath | class |
| functions | module functions |