def get_dot_graph(result: Union[Array, DictOfNamedArrays]) -> str:
r"""Return a string in the `dot <https://graphviz.org>`_ language depicting the
graph of the computation of *result*.
:arg result: Outputs of the computation (cf.
:func:`pytato.generate_loopy`).
"""
outputs: DictOfNamedArrays = normalize_outputs(result)
del result
mapper = ArrayToDotNodeInfoMapper()
for elem in outputs._data.values():
mapper(elem)
nodes = mapper.nodes
input_arrays: List[Array] = []
internal_arrays: List[ArrayOrNames] = []
array_to_id: Dict[ArrayOrNames, str] = {}
id_gen = UniqueNameGenerator()
for array in nodes:
array_to_id[array] = id_gen("array")
if isinstance(array, InputArgumentBase):
input_arrays.append(array)
else:
internal_arrays.append(array)
emit = DotEmitter()
with emit.block("digraph computation"):
emit("node [shape=rectangle]")
# Emit inputs.
with emit.block("subgraph cluster_Inputs"):
emit('label="Inputs"')
for array in input_arrays:
_emit_array(emit, nodes[array].title, nodes[array].fields,
array_to_id[array])
# Emit non-inputs.
for array in internal_arrays:
_emit_array(emit, nodes[array].title, nodes[array].fields,
array_to_id[array])
# Emit edges.
for array, node in nodes.items():
for label, tail_array in node.edges.items():
tail = array_to_id[tail_array]
head = array_to_id[array]
emit('%s -> %s [label="%s"]' % (tail, head, dot_escape(label)))
# Emit output/namespace name mappings.
_emit_name_cluster(emit,
outputs._data,
array_to_id,
id_gen,
label="Outputs")
return emit.get()
def get_num_nodes(outputs: Union[Array, DictOfNamedArrays]) -> int:
"""Returns the number of nodes in DAG *outputs*."""
from pytato.codegen import normalize_outputs
outputs = normalize_outputs(outputs)
ncm = NodeCountMapper()
ncm(outputs)
return ncm.count
def get_nusers(
outputs: Union[Array, DictOfNamedArrays]) -> Mapping[Array, int]:
"""
For the DAG *outputs*, returns the mapping from each node to the number of
nodes using its value within the DAG given by *outputs*.
"""
from pytato.codegen import normalize_outputs
outputs = normalize_outputs(outputs)
nuser_collector = NUserCollector()
nuser_collector(outputs)
return nuser_collector.nusers
def generate_loopy(result: Union[Array, DictOfNamedArrays, Dict[str, Array]],
target: Optional[LoopyTarget] = None,
options: Optional[lp.Options] = None) -> BoundProgram:
r"""Code generation entry point.
:param result: Outputs of the computation.
:param target: Code generation target.
:param options: Code generation options for the kernel.
:returns: A :class:`pytato.program.BoundProgram` wrapping the generated
:mod:`loopy` program.
"""
orig_outputs: DictOfNamedArrays = normalize_outputs(result)
del result
if target is None:
target = LoopyPyOpenCLTarget()
preproc_result = preprocess(orig_outputs)
outputs = preproc_result.outputs
compute_order = preproc_result.compute_order
namespace = outputs.namespace
state = get_initial_codegen_state(namespace, target, options)
# Reserve names of input and output arguments.
for val in namespace.values():
if isinstance(val, InputArgumentBase):
state.var_name_gen.add_name(val.name)
state.var_name_gen.add_names(outputs)
# Generate code for graph nodes.
cg_mapper = CodeGenMapper()
for name, val in sorted(
namespace.items(),
key=lambda x: x[0] # lexicographic order of names
):
_ = cg_mapper(val, state)
# Generate code for outputs.
for name in compute_order:
expr = outputs[name]
insn_id = add_store(name, expr, cg_mapper(expr, state), state)
# replace "expr" with the created stored variable
state.results[expr] = StoredResult(name, expr.ndim,
frozenset([insn_id]))
return target.bind_program(program=state.program,
bound_arguments=preproc_result.bound_arguments)
def generate_loopy(
result: Union[Array, DictOfNamedArrays, Dict[str, Array]],
target: Optional[LoopyTarget] = None,
options: Optional[lp.Options] = None,
*,
cl_device: Optional["pyopencl.Device"] = None,
array_tag_t_to_not_propagate: FrozenSet[Type[Tag]] = frozenset(
[ImplStored, Named, PrefixNamed]),
axis_tag_t_to_not_propagate: FrozenSet[Type[Tag]] = frozenset(),
) -> BoundProgram:
r"""Code generation entry point.
:param result: Outputs of the computation.
:param target: Code generation target.
:param options: Code generation options for the kernel.
:returns: A :class:`pytato.target.BoundProgram` wrapping the generated
:mod:`loopy` program.
If *result* is a :class:`dict` or a :class:`pytato.DictOfNamedArrays` and
*options* is not supplied, then the Loopy option
:attr:`~loopy.Options.return_dict` will be set to *True*. If it is supplied,
:attr:`~loopy.Options.return_dict` must already be set to *True*.
.. note::
:mod:`pytato` metadata :math:`\mapsto` :mod:`loopy` metadata semantics:
- Inames that index over an :class:`~pytato.array.Array`'s axis in the
allocation instruction are tagged with the corresponding
:class:`~pytato.array.Axis`'s tags. The caller may choose to not
propagate axis tags of type *axis_tag_t_to_not_propagate*.
- :attr:`pytato.Array.tags` of inputs/outputs in *outputs*
would be copied over to the tags of the corresponding
:class:`loopy.ArrayArg`. The caller may choose to not
propagate array tags of type *array_tag_t_to_not_propagate*.
- Arrays tagged with :class:`pytato.tags.ImplStored` would have their
tags copied over to the tags of corresponding
:class:`loopy.TemporaryVariable`. The caller may choose to not
propagate array tags of type *array_tag_t_to_not_propagate*.
"""
result_is_dict = isinstance(result, (dict, DictOfNamedArrays))
orig_outputs: DictOfNamedArrays = normalize_outputs(result)
del result
if target is None:
target = LoopyPyOpenCLTarget(device=cl_device)
else:
if cl_device is not None:
raise TypeError("may not pass both 'target' and 'cl_device'")
preproc_result = preprocess(orig_outputs, target)
outputs = preproc_result.outputs
compute_order = preproc_result.compute_order
if options is None:
options = lp.Options(return_dict=result_is_dict)
elif isinstance(options, dict):
from warnings import warn
warn(
"Passing a dict for options is deprecated and will stop working in "
"2022. Pass an actual loopy.Options object instead.",
DeprecationWarning,
stacklevel=2)
options = lp.Options(**options)
if options.return_dict != result_is_dict:
raise ValueError("options.result_is_dict is expected to match "
"whether the returned value is a dictionary")
state = get_initial_codegen_state(target, options)
from pytato.transform import InputGatherer
ing = InputGatherer()
state.var_name_gen.add_names({
input_expr.name
for name in compute_order for input_expr in ing(outputs[name].expr)
if isinstance(input_expr, (Placeholder, SizeParam, DataWrapper))
if input_expr.name is not None
})
state.var_name_gen.add_names(outputs)
cg_mapper = CodeGenMapper(array_tag_t_to_not_propagate,
axis_tag_t_to_not_propagate)
# Generate code for outputs.
for name in compute_order:
expr = outputs[name].expr
insn_id = add_store(name, expr, cg_mapper(expr, state), state,
cg_mapper)
# replace "expr" with the created stored variable
state.results[expr] = StoredResult(name, expr.ndim,
frozenset([insn_id]))
# Why call make_reduction_inames_unique?
# Consider pt.generate_loopy(pt.sum(x) + pt.sum(x)), the generated program
# would be a single instruction with rhs: `_pt_subst() + _pt_subst()`.
# The result of pt.sum(x) is cached => same instance of InlinedResult is
# emitted for both invocations and we would be required to avoid such
# reduction iname collisions.
program = lp.make_reduction_inames_unique(state.program)
return target.bind_program(program=program,
bound_arguments=preproc_result.bound_arguments)
def get_ascii_graph(result: Union[Array, DictOfNamedArrays],
use_color: bool = True) -> str:
"""Return a string representing the computation of *result*
using the `asciidag <https://pypi.org/project/asciidag/>`_ package.
:arg result: Outputs of the computation (cf.
:func:`pytato.generate_loopy`).
:arg use_color: Colorized output
"""
outputs: DictOfNamedArrays = normalize_outputs(result)
del result
mapper = ArrayToDotNodeInfoMapper()
for elem in outputs._data.values():
mapper(elem)
nodes = mapper.nodes
input_arrays: List[Array] = []
internal_arrays: List[ArrayOrNames] = []
array_to_id: Dict[ArrayOrNames, str] = {}
id_gen = UniqueNameGenerator()
for array in nodes:
array_to_id[array] = id_gen("array")
if isinstance(array, InputArgumentBase):
input_arrays.append(array)
else:
internal_arrays.append(array)
# Since 'asciidag' prints the DAG from top to bottom (ie, with the inputs
# at the bottom), we need to invert our representation of it, that is, the
# 'parents' constructor argument to Node() actually means 'children'.
from asciidag.node import Node # type: ignore[import]
asciidag_nodes: Dict[ArrayOrNames, Node] = {}
from collections import defaultdict
asciidag_edges: Dict[ArrayOrNames, List[ArrayOrNames]] = defaultdict(list)
# Reverse edge directions
for array in internal_arrays:
for _, v in nodes[array].edges.items():
asciidag_edges[v].append(array)
# Add the internal arrays in reversed order
for array in internal_arrays[::-1]:
ary_edges = [asciidag_nodes[v] for v in asciidag_edges[array]]
if array == internal_arrays[-1]:
ary_edges.append(Node("Outputs"))
asciidag_nodes[array] = Node(f"{nodes[array].title}",
parents=ary_edges)
# Add the input arrays last since they have no predecessors
for array in input_arrays:
ary_edges = [asciidag_nodes[v] for v in asciidag_edges[array]]
asciidag_nodes[array] = Node(f"{nodes[array].title}",
parents=ary_edges)
input_node = Node("Inputs",
parents=[asciidag_nodes[v] for v in input_arrays])
from asciidag.graph import Graph # type: ignore[import]
from io import StringIO
f = StringIO()
graph = Graph(fh=f, use_color=use_color)
graph.show_nodes([input_node])
# Get the graph and remove trailing whitespace
res = "\n".join([s.rstrip() for s in f.getvalue().split("\n")])
return res