compiler

compiler

TorchCompiler

Bases: AbstractCompiler[TorchCircuit]

The class responsible of handling the compilation of a symbolic circuit to a pytorch graph.

Source code in cirkit/backend/torch/compiler.py

class TorchCompiler(AbstractCompiler[TorchCircuit]):
    """The class responsible of handling the compilation of a symbolic circuit
    to a pytorch graph.
    """

    def __init__(
        self, semiring: str = "sum-product", fold: bool = False, optimize: bool = False
    ) -> None:
        super().__init__(
            CompilerLayerRegistry(DEFAULT_LAYER_COMPILATION_RULES),
            CompilerParameterRegistry(DEFAULT_PARAMETER_COMPILATION_RULES),
            CompilerInitializerRegistry(DEFAULT_INITIALIZER_COMPILATION_RULES),
            fold=fold,
            optimize=optimize,
        )

        # The semiring being used at compile time
        self._semiring: Semiring = SemiringImpl.from_name(semiring)

        # The state of the compiler
        self._state = TorchCompilerState()

        # The registries of optimization rules
        self._layer_optimization_registry = {
            "fuse": LayerOptRegistry(dict(DEFAULT_LAYER_FUSE_OPT_RULES)),
            "shatter": LayerOptRegistry(dict(DEFAULT_LAYER_SHATTER_OPT_RULES)),
        }
        self._parameter_optimization_registry = ParameterOptRegistry(
            dict(DEFAULT_PARAMETER_OPT_RULES)
        )

    def compile_pipeline(self, sc: Circuit) -> TorchCircuit:
        # Compile the circuits following the topological ordering of the pipeline.
        for sci in pipeline_topological_ordering([sc]):
            # Check if the circuit in the pipeline has already been compiled
            if self.is_compiled(sci):
                continue

            # Compile the circuit
            self._compile_circuit(sci)

        # Return the compiled circuit (i.e., the output of the circuit pipeline)
        return self.get_compiled_circuit(sc)

    @property
    def semiring(self) -> Semiring:
        return self._semiring

    @property
    def is_fold_enabled(self) -> bool:
        return self._flags["fold"]

    @property
    def is_optimize_enabled(self) -> bool:
        return self._flags["optimize"]

    @property
    def state(self) -> TorchCompilerState:
        return self._state

    def compile_layer(self, layer: Layer) -> TorchLayer:
        """Retrieve and apply the layer's compilation rule.

        Args:
            layer (Layer): Symbolic layer to compile.

        Returns:
            TorchLayer: Compiled layer.
        """
        signature = type(layer)
        rule = self.retrieve_layer_rule(signature)
        return cast(TorchLayer, rule(self, layer))

    def compile_parameter(self, parameter: Parameter) -> TorchParameter:
        """Compile a symbolic parameter graph.
        This function will iterate through all the nodes of a `TorchParameter`
        and compile each one to create the compiler parameter graph.
        Args:
            parameter (Parameter): Symbolic parameter graph to compile.

        Returns:
            TorchParameter: Compiled parameter graph.
        """
        # A map from symbolic to compiled parameters
        compiled_nodes_map: dict[ParameterNode, TorchParameterNode] = {}

        # The parameter nodes, and their inputs
        nodes: list[TorchParameterNode] = []
        in_nodes: dict[TorchParameterNode, list[TorchParameterNode]] = {}

        # Compile the parameter by following the topological ordering
        for p in parameter.topological_ordering():
            # Compile the parameter node and make the connections
            compiled_p = self._compile_parameter_node(p)
            in_compiled_nodes = [compiled_nodes_map[pi] for pi in parameter.node_inputs(p)]
            in_nodes[compiled_p] = in_compiled_nodes
            compiled_nodes_map[p] = compiled_p
            nodes.append(compiled_p)

        # Build the parameter's computational graph
        outputs = [compiled_nodes_map[parameter.output]]
        return TorchParameter(nodes, in_nodes, outputs)

    def compile_initializer(self, initializer: Initializer) -> Callable[[Tensor], Tensor]:
        """Return the initialisation function corresponding to a symbolic initializer.

        Args:
            initializer (Initializer): Symbolic initializer.

        Returns:
            Callable[[Tensor], Tensor]: Initialization function.
        """
        # Retrieve the rule for the given initializer and compile it
        signature = type(initializer)
        rule = self.retrieve_initializer_rule(signature)
        return cast(Callable[[Tensor], Tensor], rule(self, initializer))

    def retrieve_layer_optimization_registry(self, kind: str) -> LayerOptRegistry:
        return self._layer_optimization_registry[kind]

    def retrieve_parameter_optimization_registry(self) -> ParameterOptRegistry:
        return self._parameter_optimization_registry

    def retrieve_layer_optimization_rule(
        self, kind: str, pattern: LayerOptPattern
    ) -> LayerOptApplyFunc:
        registry = self.retrieve_layer_optimization_registry(kind)
        return registry.retrieve_rule(pattern)

    def retrieve_parameter_optimization_rule(
        self, pattern: ParameterOptPattern
    ) -> ParameterOptApplyFunc:
        registry = self.retrieve_parameter_optimization_registry()
        return registry.retrieve_rule(pattern)

    def _compile_parameter_node(self, node: ParameterNode) -> TorchParameterNode:
        """Return the compiled parameter node corresponding to a symbolic parameter node.

        Args:
            node (ParameterNode): Symbolic parameter node.

        Returns:
            TorchParameterNode: Compiled parameter node.
        """
        signature = type(node)
        rule = self.retrieve_parameter_rule(signature)
        return cast(TorchParameterNode, rule(self, node))

    def _compile_circuit(self, sc: Circuit) -> TorchCircuit:
        """Compile a symbolic circuit to Torch using the compiler's parameters
        In the TorchCompiler, it is possible to enable optimizations and folding
        which are applied on the compiled circuit

        Args:
            sc (Circuit): Symbolic circuit to compile

        Returns:
            TorchCircuit: Compiled circuit with optionnal optimizations and folding
        """
        # A map from symbolic to compiled layers
        compiled_layers_map: dict[Layer, TorchLayer] = {}

        # The inputs of each layer
        in_layers: dict[TorchLayer, list[TorchLayer]] = {}

        # Compile layers by following the topological ordering
        for sl in sc.topological_ordering():
            # Compile the layer, for any layer types
            layer = self.compile_layer(sl)

            # Build the connectivity between compiled layers
            ins = [compiled_layers_map[sli] for sli in sc.layer_inputs(sl)]
            in_layers[layer] = ins
            compiled_layers_map[sl] = layer

        # Construct the sequence of output layers
        outputs = [compiled_layers_map[sl] for sl in sc.outputs]

        # Construct the tensorized circuit
        layers = list(compiled_layers_map.values())
        cc = TorchCircuit(
            sc.scope,
            layers=layers,
            in_layers=in_layers,
            outputs=outputs,
            properties=sc.properties,
        )

        # Post-process the compiled circuit, i.e.,
        # optionally apply optimizations to it and then fold it
        cc = self._post_process_circuit(cc)

        # Allocate & initialize the parameters
        cc.reset_parameters()

        # Register the compiled circuit
        self.register_compiled_circuit(sc, cc)

        # Signal the end of the circuit compilation to the state
        self._state.finish_compilation()
        return cc

    def _post_process_circuit(self, cc: TorchCircuit) -> TorchCircuit:
        """Apply the post processing steps corresponding to the active flags
        This compiler can :
            - optimize the circuit through layer's fusion / splitting
            - fold the circuit's layers
        Args:
            cc (TorchCircuit): Compiled circuit to post process

        Returns:
            TorchCircuit: Post processed compiled circuit
        """
        if self.is_optimize_enabled:
            # Optimize the circuit computational graph
            opt_cc = _optimize_circuit(self, cc, max_opt_steps=5)
            del cc
            cc = opt_cc
        if self.is_fold_enabled:
            # Optimize the circuit by folding it
            opt_cc = _fold_circuit(self, cc)
            del cc
            cc = opt_cc
        return cc

is_fold_enabled property

is_optimize_enabled property

semiring property

state property

__init__(semiring='sum-product', fold=False, optimize=False)

Source code in cirkit/backend/torch/compiler.py

def __init__(
    self, semiring: str = "sum-product", fold: bool = False, optimize: bool = False
) -> None:
    super().__init__(
        CompilerLayerRegistry(DEFAULT_LAYER_COMPILATION_RULES),
        CompilerParameterRegistry(DEFAULT_PARAMETER_COMPILATION_RULES),
        CompilerInitializerRegistry(DEFAULT_INITIALIZER_COMPILATION_RULES),
        fold=fold,
        optimize=optimize,
    )

    # The semiring being used at compile time
    self._semiring: Semiring = SemiringImpl.from_name(semiring)

    # The state of the compiler
    self._state = TorchCompilerState()

    # The registries of optimization rules
    self._layer_optimization_registry = {
        "fuse": LayerOptRegistry(dict(DEFAULT_LAYER_FUSE_OPT_RULES)),
        "shatter": LayerOptRegistry(dict(DEFAULT_LAYER_SHATTER_OPT_RULES)),
    }
    self._parameter_optimization_registry = ParameterOptRegistry(
        dict(DEFAULT_PARAMETER_OPT_RULES)
    )

compile_initializer(initializer)

Return the initialisation function corresponding to a symbolic initializer.

Parameters:

Name	Type	Description	Default
initializer	Initializer	Symbolic initializer.	required

Returns:

Type	Description
Callable[[Tensor], Tensor]	Callable[[Tensor], Tensor]: Initialization function.

Source code in cirkit/backend/torch/compiler.py

def compile_initializer(self, initializer: Initializer) -> Callable[[Tensor], Tensor]:
    """Return the initialisation function corresponding to a symbolic initializer.

    Args:
        initializer (Initializer): Symbolic initializer.

    Returns:
        Callable[[Tensor], Tensor]: Initialization function.
    """
    # Retrieve the rule for the given initializer and compile it
    signature = type(initializer)
    rule = self.retrieve_initializer_rule(signature)
    return cast(Callable[[Tensor], Tensor], rule(self, initializer))

compile_layer(layer)

Retrieve and apply the layer's compilation rule.

Parameters:

Name	Type	Description	Default
layer	Layer	Symbolic layer to compile.	required

Returns:

Name	Type	Description
TorchLayer	TorchLayer	Compiled layer.

Source code in cirkit/backend/torch/compiler.py

def compile_layer(self, layer: Layer) -> TorchLayer:
    """Retrieve and apply the layer's compilation rule.

    Args:
        layer (Layer): Symbolic layer to compile.

    Returns:
        TorchLayer: Compiled layer.
    """
    signature = type(layer)
    rule = self.retrieve_layer_rule(signature)
    return cast(TorchLayer, rule(self, layer))

compile_parameter(parameter)

Compile a symbolic parameter graph. This function will iterate through all the nodes of a TorchParameter and compile each one to create the compiler parameter graph. Args: parameter (Parameter): Symbolic parameter graph to compile.

Returns:

Name	Type	Description
TorchParameter	TorchParameter	Compiled parameter graph.

Source code in cirkit/backend/torch/compiler.py

def compile_parameter(self, parameter: Parameter) -> TorchParameter:
    """Compile a symbolic parameter graph.
    This function will iterate through all the nodes of a `TorchParameter`
    and compile each one to create the compiler parameter graph.
    Args:
        parameter (Parameter): Symbolic parameter graph to compile.

    Returns:
        TorchParameter: Compiled parameter graph.
    """
    # A map from symbolic to compiled parameters
    compiled_nodes_map: dict[ParameterNode, TorchParameterNode] = {}

    # The parameter nodes, and their inputs
    nodes: list[TorchParameterNode] = []
    in_nodes: dict[TorchParameterNode, list[TorchParameterNode]] = {}

    # Compile the parameter by following the topological ordering
    for p in parameter.topological_ordering():
        # Compile the parameter node and make the connections
        compiled_p = self._compile_parameter_node(p)
        in_compiled_nodes = [compiled_nodes_map[pi] for pi in parameter.node_inputs(p)]
        in_nodes[compiled_p] = in_compiled_nodes
        compiled_nodes_map[p] = compiled_p
        nodes.append(compiled_p)

    # Build the parameter's computational graph
    outputs = [compiled_nodes_map[parameter.output]]
    return TorchParameter(nodes, in_nodes, outputs)

compile_pipeline(sc)

Source code in cirkit/backend/torch/compiler.py

def compile_pipeline(self, sc: Circuit) -> TorchCircuit:
    # Compile the circuits following the topological ordering of the pipeline.
    for sci in pipeline_topological_ordering([sc]):
        # Check if the circuit in the pipeline has already been compiled
        if self.is_compiled(sci):
            continue

        # Compile the circuit
        self._compile_circuit(sci)

    # Return the compiled circuit (i.e., the output of the circuit pipeline)
    return self.get_compiled_circuit(sc)

retrieve_layer_optimization_registry(kind)

Source code in cirkit/backend/torch/compiler.py

def retrieve_layer_optimization_registry(self, kind: str) -> LayerOptRegistry:
    return self._layer_optimization_registry[kind]

retrieve_layer_optimization_rule(kind, pattern)

Source code in cirkit/backend/torch/compiler.py

def retrieve_layer_optimization_rule(
    self, kind: str, pattern: LayerOptPattern
) -> LayerOptApplyFunc:
    registry = self.retrieve_layer_optimization_registry(kind)
    return registry.retrieve_rule(pattern)

retrieve_parameter_optimization_registry()

Source code in cirkit/backend/torch/compiler.py

def retrieve_parameter_optimization_registry(self) -> ParameterOptRegistry:
    return self._parameter_optimization_registry

retrieve_parameter_optimization_rule(pattern)

Source code in cirkit/backend/torch/compiler.py

def retrieve_parameter_optimization_rule(
    self, pattern: ParameterOptPattern
) -> ParameterOptApplyFunc:
    registry = self.retrieve_parameter_optimization_registry()
    return registry.retrieve_rule(pattern)

TorchCompilerState

Source code in cirkit/backend/torch/compiler.py

class TorchCompilerState:
    def __init__(self) -> None:
        # A map from symbolic parameter tensors to a tuple containing the compiled parameter tensor,
        # and the slice index, which is 0 if the compiled parameter tensor is unfolded.
        # If the compiled parameter tensor is folded, then the slice index can be non-zero.
        self._compiled_parameters: dict[TensorParameter, tuple[TorchTensorParameter, int]] = {}

        # We keep a reverse map from compiled and unfolded parameter tensors
        # to the corresponding symbolic parameter tensors.
        # This is useful to update the map from symbolic to compiled parameter tensors above
        # after we fold the tensor parameters within a circuit.
        # Since this is useful only for folding, it will be cleared after each circuit compilation.
        self._symbolic_parameters: dict[TorchTensorParameter, TensorParameter] = {}

    def finish_compilation(self) -> None:
        # Clear the map from (unfolded) compiled parameter tensors to symbolic ones
        self._symbolic_parameters.clear()

    def has_compiled_parameter(self, p: TensorParameter) -> bool:
        # Retrieve whether a tensor parameter has already been compiled
        return p in self._compiled_parameters

    def retrieve_compiled_parameter(self, p: TensorParameter) -> tuple[TorchTensorParameter, int]:
        # Retrieve the compiled parameter: we return the fold index as well.
        return self._compiled_parameters[p]

    def retrieve_symbolic_parameter(self, p: TorchTensorParameter) -> TensorParameter:
        # Retrieve the symbolic parameter tensor associated to the compiled one (which is unfolded)
        return self._symbolic_parameters[p]

    def register_compiled_parameter(
        self,
        sp: TensorParameter,
        cp: TorchTensorParameter,
        *,
        fold_idx: int | None = None,
    ) -> None:
        # Register a link from a symbolic parameter tensor to a compiled parameter tensor.
        if fold_idx is None:
            # We are registering an unfolded compiled parameter tensor
            # So, we can also register the reverse map (i.e., compiled to symbolic)
            self._compiled_parameters[sp] = (cp, 0)
            self._symbolic_parameters[cp] = sp
        else:
            # We are registering a folded compiled parameter tensor
            # So, we associate the symbolic parameter tensor to a particular slice of the
            # folded compiled parameter tensor, which is specified by the 'fold_idx'.
            self._compiled_parameters[sp] = (cp, fold_idx)

__init__()

Source code in cirkit/backend/torch/compiler.py

def __init__(self) -> None:
    # A map from symbolic parameter tensors to a tuple containing the compiled parameter tensor,
    # and the slice index, which is 0 if the compiled parameter tensor is unfolded.
    # If the compiled parameter tensor is folded, then the slice index can be non-zero.
    self._compiled_parameters: dict[TensorParameter, tuple[TorchTensorParameter, int]] = {}

    # We keep a reverse map from compiled and unfolded parameter tensors
    # to the corresponding symbolic parameter tensors.
    # This is useful to update the map from symbolic to compiled parameter tensors above
    # after we fold the tensor parameters within a circuit.
    # Since this is useful only for folding, it will be cleared after each circuit compilation.
    self._symbolic_parameters: dict[TorchTensorParameter, TensorParameter] = {}

finish_compilation()

Source code in cirkit/backend/torch/compiler.py

def finish_compilation(self) -> None:
    # Clear the map from (unfolded) compiled parameter tensors to symbolic ones
    self._symbolic_parameters.clear()

has_compiled_parameter(p)

Source code in cirkit/backend/torch/compiler.py

def has_compiled_parameter(self, p: TensorParameter) -> bool:
    # Retrieve whether a tensor parameter has already been compiled
    return p in self._compiled_parameters

register_compiled_parameter(sp, cp, *, fold_idx=None)

Source code in cirkit/backend/torch/compiler.py

def register_compiled_parameter(
    self,
    sp: TensorParameter,
    cp: TorchTensorParameter,
    *,
    fold_idx: int | None = None,
) -> None:
    # Register a link from a symbolic parameter tensor to a compiled parameter tensor.
    if fold_idx is None:
        # We are registering an unfolded compiled parameter tensor
        # So, we can also register the reverse map (i.e., compiled to symbolic)
        self._compiled_parameters[sp] = (cp, 0)
        self._symbolic_parameters[cp] = sp
    else:
        # We are registering a folded compiled parameter tensor
        # So, we associate the symbolic parameter tensor to a particular slice of the
        # folded compiled parameter tensor, which is specified by the 'fold_idx'.
        self._compiled_parameters[sp] = (cp, fold_idx)

retrieve_compiled_parameter(p)

Source code in cirkit/backend/torch/compiler.py

def retrieve_compiled_parameter(self, p: TensorParameter) -> tuple[TorchTensorParameter, int]:
    # Retrieve the compiled parameter: we return the fold index as well.
    return self._compiled_parameters[p]

retrieve_symbolic_parameter(p)

Source code in cirkit/backend/torch/compiler.py

def retrieve_symbolic_parameter(self, p: TorchTensorParameter) -> TensorParameter:
    # Retrieve the symbolic parameter tensor associated to the compiled one (which is unfolded)
    return self._symbolic_parameters[p]