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