"""Module for compilation artifacts.""" import inspect import platform import shutil import subprocess from pathlib import Path from typing import Any, Callable, Dict, Optional, Union import networkx as nx from PIL import Image from ..debugging import draw_graph, get_printable_graph from ..operator_graph import OPGraph from ..representation import intermediate as ir from ..values import BaseValue DEFAULT_OUTPUT_DIRECTORY: Path = Path(".artifacts") class CompilationArtifacts: """Class that conveys information about compilation process.""" output_directory: Path source_code_of_the_function_to_compile: Optional[str] parameters_of_the_function_to_compile: Dict[str, str] drawings_of_operation_graphs: Dict[str, Image.Image] textual_representations_of_operation_graphs: Dict[str, str] final_operation_graph: Optional[OPGraph] bounds_of_the_final_operation_graph: Optional[Dict[ir.IntermediateNode, Dict[str, Any]]] mlir_of_the_final_operation_graph: Optional[str] def __init__(self, output_directory: Path = DEFAULT_OUTPUT_DIRECTORY): self.output_directory = output_directory self.source_code_of_the_function_to_compile = None self.parameters_of_the_function_to_compile = {} self.drawings_of_operation_graphs = {} self.textual_representations_of_operation_graphs = {} self.final_operation_graph = None self.bounds_of_the_final_operation_graph = None self.mlir_of_the_final_operation_graph = None def add_function_to_compile(self, function: Union[Callable, str]): """Adds the function to compile to artifacts. Args: function (Union[Callable, str]): the function to compile or source code of it Returns: None """ self.source_code_of_the_function_to_compile = ( function if isinstance(function, str) else inspect.getsource(function) ) def add_parameter_of_function_to_compile(self, name: str, value: Union[BaseValue, str]): """Adds a parameter of the function to compile to the artifacts. Args: name (str): name of the parameter value (Union[BaseValue, str]): value of the parameter or textual representation of it Returns: None """ self.parameters_of_the_function_to_compile[name] = str(value) def add_operation_graph(self, name: str, operation_graph: OPGraph): """Adds an operation graph to the artifacts. Args: name (str): name of the graph operation_graph (OPGraph): the operation graph itself Returns: None """ drawing = draw_graph(operation_graph) textual_representation = get_printable_graph(operation_graph, show_data_types=True) self.drawings_of_operation_graphs[name] = drawing self.textual_representations_of_operation_graphs[name] = textual_representation self.final_operation_graph = operation_graph def add_final_operation_graph_bounds(self, bounds: Dict[ir.IntermediateNode, Dict[str, Any]]): """Adds the bounds of the final operation graph to the artifacts. Args: bounds (Dict[ir.IntermediateNode, Dict[str, Any]]): the bound dictionary Returns: None """ assert self.final_operation_graph is not None self.bounds_of_the_final_operation_graph = bounds def add_final_operation_graph_mlir(self, mlir: str): """Adds the mlir of the final operation graph to the artifacts. Args: mlir (str): the mlir code of the final operation graph Returns: None """ assert self.final_operation_graph is not None self.mlir_of_the_final_operation_graph = mlir def export(self): """Exports the artifacts to a the output directory. Returns: None """ output_directory = self.output_directory if output_directory.exists(): shutil.rmtree(output_directory) output_directory.mkdir() with open(output_directory.joinpath("environment.txt"), "w") as f: f.write(f"{platform.platform()} {platform.version()}\n") f.write(f"Python {platform.python_version()}\n") with open(output_directory.joinpath("requirements.txt"), "w") as f: # example `pip list` output # Package Version # ----------------------------- --------- # alabaster 0.7.12 # appdirs 1.4.4 # ... ... # ... ... # wrapt 1.12.1 # zipp 3.5.0 pip_process = subprocess.run(["pip", "list"], stdout=subprocess.PIPE, check=True) dependencies = iter(pip_process.stdout.decode("utf-8").split("\n")) # skip 'Package ... Version' line next(dependencies) # skip '------- ... -------' line next(dependencies) for dependency in dependencies: tokens = [token for token in dependency.split(" ") if token != ""] if len(tokens) == 0: continue name = tokens[0] version = tokens[1] f.write(f"{name}=={version}\n") if self.source_code_of_the_function_to_compile is not None: with open(output_directory.joinpath("function.txt"), "w") as f: f.write(self.source_code_of_the_function_to_compile) if len(self.parameters_of_the_function_to_compile) > 0: with open(output_directory.joinpath("parameters.txt"), "w") as f: for name, parameter in self.parameters_of_the_function_to_compile.items(): f.write(f"{name} :: {parameter}\n") drawings = self.drawings_of_operation_graphs.items() for index, (name, drawing) in enumerate(drawings): identifier = CompilationArtifacts._identifier(index, name) drawing.save(output_directory.joinpath(f"{identifier}.png")) textual_representations = self.textual_representations_of_operation_graphs.items() for index, (name, representation) in enumerate(textual_representations): identifier = CompilationArtifacts._identifier(index, name) with open(output_directory.joinpath(f"{identifier}.txt"), "w") as f: f.write(f"{representation}\n") if self.bounds_of_the_final_operation_graph is not None: assert self.final_operation_graph is not None with open(output_directory.joinpath("bounds.txt"), "w") as f: # TODO: # if nx.topological_sort is not deterministic between calls, # the lines below will not work properly # thus, we may want to change this in the future for index, node in enumerate(nx.topological_sort(self.final_operation_graph.graph)): bounds = self.bounds_of_the_final_operation_graph.get(node) assert bounds is not None f.write(f"%{index} :: [{bounds.get('min')}, {bounds.get('max')}]\n") if self.mlir_of_the_final_operation_graph is not None: assert self.final_operation_graph is not None with open(output_directory.joinpath("mlir.txt"), "w") as f: f.write(self.mlir_of_the_final_operation_graph) @staticmethod def _identifier(index, name): return f"{index + 1}.{name}.graph"