# ShortCircuit is a project for generating an AIG out of a truth-table

The project is organized according to the template below

```
│
├── conf                      <- Hydra configs
│   ├── data                     <- Data configs
│   ├── experiment               <- Experiment configs
│   ├── model                    <- Model configs
│   ├── loss                     <- Loss configs
│   ├── optimizer                <- Optimizer configs
│   ├── scheduler                <- Scheduler configs
│   ├── policy                   <- Policy configs
│   │
│   ├── eval.yaml             <- Main config for evaluation
│   ├── train.yaml            <- Main config for supervised pre-training
│   ├── finetune.yaml         <- Main config for RL fine-tuning
│   └── utils.yaml            <- Main config for utilities 
│
├── data                   <- Project data
│
├── logs                   <- Logs generated by hydra and pytorch loggers
│
├── src                    <- Source code
│   ├── data                     <- Data scripts
│   │   ├── abc                       <- State-of-the-art logic synthesis tool
│   │   └── pyaig                     <- A library for handling AIGs and truth-tables
│   │
│   ├── loss                     <- Loss scripts
│   ├── model                    <- Model scripts
│   ├── optimizer                <- Optimizer scripts
│   ├── scheduler                <- Scheduler scripts
│   ├── utils                    <- Utility scripts
│   ├── search                   <- Utility scripts
│   │
│   ├── eval.py                  <- Run evaluation
│   └── train.py                 <- Run training
│
└── README.md
```

## Install Dependencies
To install the dependencies for this project run:
```shell
pip install -r requirements.txt
```
In case any dependendamcy is missing please install it when prompted. Also, you should use a conda env. 


## Hydra and Configs
We use hydra to keep track of experiments and hyperparameters. To get the available settings for a script you can run:
```shell
python src/[script].py --help
```


## Download Datasets
To generate any data for this project large AIGs are necessary. In the paper we use the EPFL benchmarks and download them with:
```shell
git clone git@github.com:lsils/benchmarks.git
```


## Generate Pre-Training Data
To generate pretraining data of [x]-input AIGs you can use the command, and providing the path to the EPFL benchmark directory:
```shell
python src/utils.py +f=extract_multicuts +data_path=[path to directory]/**/*.aig ++num_inputs=[x]
```
This will create a directory under `data/unoptimized/[x]-inputs` with many `[name].aig` files.


## Optimize Pre-Training Data with ABC
To optimize AIGs with ABC you have to build abc first which we have included in src/data/abc. Please follow the instructions in their README.md file to compile ABC. To optimize the AIGs you can use the command:
```shell
python src/utils.py +f=optimize_cuts +data_path=data/unoptimized/[x]-inputs/**/*.aig
```
This will create a directory under `data/optimized/[x]-inputs` with many `[name].aig` files.


## Extract Truth Tables for Fine-Tuning
To extract truth tables out of the training data you can run the command:
```shell
python src/utils.py +f=extract_truth_tables +data_path=data/unoptimized/[x]-inputs/**/*.aig
```
This will create a file with path: `data/truth_tables/extracted/[x]-inputs.txt`.


## Mlflow Server Setup
Unfortunately, the code has not been cleaned to only log in the terminal. We use a private mlflow server where we log all the metadata produced during training and evaluation results. To install mlflow run:
```shell
pip install mlflow
```
Then to launch the mlflow server run:
```shell
mlflow server --host localhost --polrt 5000
```
For the rest of the experiments please set the `mlflow_uri` to `http://localhost:5000`. Once you launch the either training or evaluation you can open the link in your browser to view all the data.


## Configs
We provide some sample configuration files:
    - For pre-training: `conf/experiment/8_input_pretraining.yaml`
    - For fine-tuning: `conf/experiment/8_input_finetuning.yaml`
    - For evaluation: `conf/experiment/8_input_eval.yaml`


## Supervised Pre-Training
To run training use the script `/src/train.py`. We use a few GPUs to train our model so you need to use `torchrun`. Before the script starts the training, it starts to extract the truth table and AIG metadata and stores in some Memmory Mapped Tensor i.e. Tensors that live on the distk. This process is lenghthy and can take about 1hr. Subsequent runs will use the cache and load the necessary meta-data in a few minutes. While this is an issue we plan to fix later on, please launch only 1 process the first time to build the cache, kill it once the cache is built and then run the training with the desired number of processes. Also, we assume that the main process that logs information and saves the model is on cuda:0. If not please kill your colleagues' training and start the training with cuda:0 available.
```shell
torchrun --standalone src/train.py +experiment=[your conf e.g. 8_input_pretraining]
```
The produced model can be used for evaluation, but need to disable the value network.


## RL Fine-Tuning
To train either from scratch or finetune a model with AlphaZero use the script `/src/finetune.py`. Unfortunately, to the distribute RL training you need to build a specific TorchRL branch from source. To build torchrl you need to run:
```shell
git clone git@github.com:pytorch/rl.git
cd rl
git checkout example-rb-rpc # The specific branch
python setup.py develop
```
Common pitfals: make sure that you unintall torchrl before installing it from source.
With TorchRL installed you can launch finetuning.
```shell
python src/finetune.py +experiment=[your conf e.g. 8_input_finetuning]
```


## Evaluation for Pre-Trained/Fine-Tuned Models
To evaluate the quality of the model use the script `/src/eval.py`. To see the available options and default configurations run:
```shell
python src/eval.py +experiment=[your conf e.g. 8_input_eval]
```


## Overriding Parameters
In case you would like to override any of the parameters from the CLI you can append the command with `++parameter=value`.
For example to override and set the epochs to 10 we can run:
```shell
python src/[script].py +experiment=[conf] ++epochs=10
```

