Compilation

Circuit Compilation Algorithm (TorchCompiler._compile_circuit)

Data structures:

• compiled_layers_map: dictionary storing mapping between symbolic layers and compiled layers.
• in_layers: dictionary storing the links between layers (adjacency list). That is, for each layer, store a list of input nodes.

Algorithm :

1. For every layer in topological order:
2. Call the compile_layer method.
3. Store the mapping between the compiled and symbolic layers in compiled_layers_map.
4. Store the mapping between the new compiled layer and its compiled inputs in the in_layers dictionary. Due to the topological order, the inputs are guaranteed to be already stored in compiled_layers_map
5. Once all the layers are compiled, create the sequence outputs which stores the compiled output layers. To do so, you just retrieve the compiled layers corresponding to the symbolic circuit's output layers.
6. Create the fully compiled circuit, a TorchCircuit object, using the mappings and sequence previously built.
7. Run the circuit post processing function TorchCompiler._post_process_circuit (optimization and folding).
8. Initialize all the parameters of the circuit (TorchCircuit.reset_parameters). It's a function that recursively calls the compiled initializer of all the circuit's parameters.
9. Register the compiled circuit in the CompiledCircuitsMap.

Layer Compilation Process (TorchCompiler.compile_layer)

The function, in itself, only retrieves the layer rule corresponding to the input type from the CompilerLayerRegistry and applies it. Here is a flowchart describing the process of compiling a single layer:

flowchart TD
    A["Symbolic Layer"] -- Layer Type --> LReg("LayerRegistry")
    LReg -- Call Layer Rule --> hasP{"Has Parameters"}
    hasP -- No --> CL["Compiled Layer"]
    hasP -- Yes --> PComp("Compile TorchParameter Graph")

    subgraph Compile Parameter
      PComp --> nextPara("Next parameter node")

      nextPara -- Parameter Type --> PReg("ParameterRegistry")
      PReg --  Call Parameter Rule --> hasI{"Has Initializer"}
      hasI -- No --> CL
      hasI -- Yes --> IComp("Compile Initializer")
      IComp -- Initializer Type --> IReg("InitializerRegistry")
      IReg -- Call Initializer Rule --> ILParam{"Is last Parameter Node ?"}
      ILParam -- Yes --> CL
      ILParam -- No --> nextPara
    end

As we can see, the compilation rule of a layer might call the compilation rule of a parameter which might itself call the compilation rule of an initializer. Furthermore, as TorchParameter objects are graphs themselves, we need to iterate through all parameter nodes to compile them.

Compile a Parameter Graph (TorchCompiler.compile_parameter)

Data structures:

• compiled_nodes_map: dictionary storing mapping between symbolic parameter nodes and compiled parameter nodes.
• in_nodes: dictionary storing the links between nodes. That is, for each nodes, store a list of input nodes.
• nodes: list of compiled nodes.

Algorithm:

1. For every parameter node in topological order:
2. Call the _compile_parameter_node method.
3. Store the mapping between the compiled and symbolic node in compiled_nodes_map.
4. Store the mapping between the new compiled node and its compiled inputs in in_nodes. Due to the topological order, the inputs are guaranteed to be already stored in compiled_nodes_map
5. Store the compiled node in nodes.
6. Once all the nodes are compiled, create the sequence outputs which stores the compiled output nodes. To do so, you simply retrieve the compiled nodes corresponding to the parameter graph output nodes.
7. Create the fully compiled graph, a TorchParameter object, using the mappings and sequence previously built.