experiments training logs; evaluation plots; update README

.gitignore

@@ -5,7 +5,7 @@ experiment_dir
 mlruns
 dataset
 
-*.ipynb
+# *.ipynb
 *.txt
 # *.json
 *.sh

README.md

@@ -1,10 +1,43 @@
-# Workflow
+# Context
+The project aims to perform Automatic Code Optimization using reinforcement learning in the Tiramisu compiler. We investigate the use of pretrained autoencoders to get the program's access matrices embeddings, which reduces the size of computational vectors.
+
+# Project Setup
+## 1. Installing Tiramisu
+Here is a detailed and updated [guide](https://tremendous-radio-85b.notion.site/Tiramisu-Installation-Guide-3f5eed16d5d641bb8ed45c493ca2ca8e) on installing Tiramisu compiler.
+
+## 2. Running RL Agent
+Here is a detailed [guide](https://kamel-brouthen.notion.site/Running-RL-Agent-07df4e9c793b44e7a5d30ed057f2b079https://kamel-brouthen.notion.site/Running-RL-Agent-07df4e9c793b44e7a5d30ed057f2b079) on running the RL agent. Additionally, the guide includes steps to set up Cuda and PyTorch for GPU training.
+
+## 3. Codebase
+The project is based on the following work:
+
+### 3.1 Reinforcement Learning Agent applying Proximal Policy Optimization (PPO) with a Graph Neural Network (GNN) backbone
+- Repository: [gnn_rl](https://github.com/Tiramisu-Compiler/gnn_rl) by [Lamouri Djamel Rassem](https://github.com/djamelrassem)
+
+### 3.2 Pretrained Autoencoder for Tiramisu Program's Computational Vector
+- Repository: [cost_model_pretrain](https://github.com/Tiramisu-Compiler/cost_model_pretrain) by Chunting Liu
+
+## 4. Training Workflow
+- Create conda environment
 ```shell
 conda env create -f environment.yaml
 ```
+- Activate conda environment
 ```shell
 conda activate tiramisu-build-env
 ```
+- Train the agent
 ```shell
 python train_ppo_gnn.py 
 ```
+- Evaluate the agent
+ ```shell
+python evaluate_ppo_gnn.py 
+```
+
+# 5. Resources
+- [Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code](https://arxiv.org/abs/1804.10694)
+- [A Deep Learning Based Cost Model for Automatic Code Optimization](https://arxiv.org/abs/2104.04955https://arxiv.org/abs/2104.04955)
+- [LOOPer: A Learned Automatic Code Optimizer For Polyhedral Compilers](https://arxiv.org/abs/2403.11522)
+- [Utilisation d’apprentissage par renforcement pour l’optimisation automatique de code dans TiramisuUtilisation d’apprentissage par renforcement pour l’optimisation automatique de code dans Tiramisu](https://www.researchgate.net/publication/372128690_Utilisation_d'apprentissage_par_renforcement_pour_l'optimisation_automatique_de_code_dans_Tiramisu)
+- [Automatic Code Optimization in the MLIR Compiler Using Deep Reinforcement Learning](https://www.researchgate.net/publication/382047058_Automatic_Code_Optimization_in_the_MLIR_Compiler_Using_Deep_Reinforcement_Learning)

agent/graph_utils.py

@@ -81,7 +81,6 @@ def comps_to_vectors(annotations, embed_access_matrices, embedding_type):
         device = "cuda:0" if torch.cuda.is_available() else "cpu"
     else:
         comp_vector_size = 709 + 9
-
     max_depth = 5
     dict_comp = {}
     for comp in annotations["computations"]:

config/config.yaml

@@ -36,8 +36,8 @@ experiment:
     NB_EXEC: 3
 
 pretrain:
-    embed_access_matrices: True
-    embedding_type: final_hidden_state # final_hidden_state, final_cell_state, concat_final_hidden_cell_state, mean_pooling_output, max_pooling_output, flattened_output
+    embed_access_matrices: False
+    embedding_type: concat_max_pooling_output_final_hidden_state # final_hidden_state, final_cell_state, concat_final_hidden_cell_state, mean_pooling_output, max_pooling_output, flattened_output, concat_max_pooling_output_final_hidden_state
 
 hyperparameters:
     num_updates: 500

evaluate_ppo_gnn.py

@@ -45,11 +45,13 @@ def write_cpp_file(schedule_object):
         input_size = 6 + get_embedding_size(Config.config.pretrain.embedding_type) + 9
     else:
         input_size = 718
+
+    print(input_size)
 
     ppo_agent = GAT(input_size=input_size, num_heads=4, hidden_size=128, num_outputs=56).to(
         device
     )
-    
+
     model_name = 'model_full_computational_vector_u500_b500_ent0.5_426'
     ppo_agent.load_state_dict(
         torch.load(