Merge pull request #2 from Armanddevacc/feat/asm-dataflow-encoder

README.md

@@ -1,242 +1,187 @@
 <h1 align="center">
-    <p>⚡ <b>EB-JEPA</b></p>
+    <p>⚙️ <b>JEPA-ASML</b></p>
 </h1>
 
 <h2 align="center">
-    <p><i>Energy-Based Joint-Embedding Predictive Architectures</i></p>
+    <p><i>JEPA pour la superoptimisation d'assembleur</i></p>
 </h2>
 
-<div align="center" style="line-height: 1;">
-  <a href="https://github.com/Trick5t3r/eb_jepa" target="_blank" style="margin: 2px;"><img alt="Github" src="https://img.shields.io/badge/Github-Trick5t3r/eb__jepa-black?logo=github" style="display: inline-block; vertical-align: middle;"/></a>
-  <a href="https://arxiv.org/abs/2602.03604" target="_blank" style="margin: 2px;"><img alt="ArXiv" src="https://img.shields.io/badge/arXiv-2602.03604-b5212f?logo=arxiv" style="display: inline-block; vertical-align: middle;"/></a>
-</div>
-
-<br>
-
 <p align="center">
-  <b><a href="https://ai.facebook.com/research/">Meta AI Research, FAIR</a></b>
+  Apprendre une représentation du <b>sens</b> d'un programme assembleur — invariante
+  à la forme produite par le compilateur — et un <b>prédicteur world-model</b> qui
+  transforme la représentation d'un code lent (O0) en celle de son équivalent
+  optimisé (O3).
 </p>
 
-<p align="center">
-  <a href="https://x.com/BasileTerv987">Basile Terver</a>,
-  Randall Balestriero,
-  Megi Dervishi,
-  David Fan,
-  Quentin Garrido,
-  Tushar Nagarajan,
-  <br>
-  Koustuv Sinha,
-  Wancong Zhang,
-  Mike Rabbat,
-  Yann LeCun,
-  Amir Bar
-</p>
-
-<p align="center">
-  An open source library and tutorial for learning representations for<br>
-  prediction and planning using joint embedding predictive architectures.
-</p>
-
-<p align="center">
-  <img src="docs/archi-schema-eb-jepa.png" alt="EB-JEPA Architecture" width="800">
-</p>
-
-> Each example is (almost) self-contained and training takes up to a few hours on a single GPU card.
-
 ---
 
-## 📚 Examples
+## 🎯 Le projet en une page
 
-### [Image JEPA](examples/image_jepa/README.md)
+Un compilateur peut produire, à partir d'une même fonction, des assembleurs très
+différents selon le niveau d'optimisation (`-O0` … `-O3`). Ces variantes ont le
+**même sens** (mêmes entrées → mêmes sorties) mais des **formes** et des **coûts**
+(vitesse) très différents.
 
-Self-supervised representations from unlabeled images on CIFAR-10, evaluated on classification.
+L'idée de ce projet est d'appliquer les **Joint-Embedding Predictive
+Architectures** (JEPA) à ce problème, avec une représentation *factorisée* :
 
-![Image JEPA Architecture](examples/image_jepa/assets/arch_figure.png)
+| Facteur | Ce qu'il capture | Doit être… |
+|---|---|---|
+| **meaning** (sens) | la sémantique du bloc | invariant à l'optimisation (O0 ≈ O3) |
+| **temporality** (coût) | la vitesse / le coût d'exécution | sensible à l'optimisation |
+| **architecture** | la cible matérielle | conditionnement (travaux futurs) |
 
-### [Video JEPA](examples/video_jepa/README.md)
+Une fois cet espace appris, la **superoptimisation** devient un problème de
+*planning* dans l'espace latent : partir de `repr(O0)`, appliquer une **action**
+(« optimise vers O3 ») et atteindre `repr(O3)`, puis décoder vers de l'assembleur.
 
-Predict next image representation in a sequence.
+> Ce dépôt est un **fork** de la librairie [EB-JEPA](#-crédits--upstream) de Meta
+> AI (FAIR). Tout le code propre au projet vit dans `eb_jepa/asm/` et
+> `examples/asm_superopt/` ; le reste est la librairie JEPA d'origine, réutilisée
+> telle quelle.
 
-![Moving MNIST](examples/video_jepa/assets/viz.png)
+---
 
-### [AC Video JEPA](examples/ac_video_jepa/README.md)
+## 🧩 Le pipeline
 
-JEPA for world modeling + planning in Two Rooms environment.
+```
+  fichier .c
+     │  clang -S -O{0,1,2,3} -masm=intel
+     ▼
+  assembleur x86-64 (Intel)                     eb_jepa/asm/corpus.py
+     │  découpage en blocs de base (straight-line, sans branchements)
+     ▼
+  graphe de flot de données (DataFlowGraph)     eb_jepa/asm/dataflow.py
+     │  nœuds = instructions ; arêtes = dépendances def→use (REG / MEM / FLAGS)
+     ▼
+  encodeur GNN (message passing typé + mean-pool)   eb_jepa/asm/encoder.py
+     │  têtes factorisées → vecteur « meaning » + sortie « temporality »
+     ▼
+  prédicteur world-model :  g(O0) + action(→O3)  ≈  g(O3)
+```
 
-| Planning Episode | Task Definition |
-|------------------|-----------------|
-| <img src="examples/ac_video_jepa/assets/top_randw_agent_steps_succ.gif" alt="Successful planning episode" width="155" /> | <img src="examples/ac_video_jepa/assets/top_randw_state.png" alt="Episode task definition" width="300" /> |
-| *Successful planning episode* | *From init to goal state* |
+Le **graphe de flot de données** expose volontairement l'*indépendance* entre
+instructions (deux instructions sans chemin def→use n'ont pas d'arête) — c'est
+exactement ce qui rend détectable le parallélisme de bloc (packing SLP / AVX) en
+aval. Aucun code n'est exécuté ni assemblé : tout part du texte assembleur.
 
 ---
 
-## 🚀 Installation
-
-### HTW cluster — quick start (hackathon only)
+## 🏋️ Entraînement en deux phases
 
-> Skip this section unless you are on the HTW hackathon cluster — the generic install below is all you need locally.
+Le point clé du design : l'encodeur n'utilise **jamais** les labels de niveau
+d'optimisation. Le couplage O0↔O3 est réservé au prédicteur.
 
-Please follow the [setup instructions](setup.md) before starting the project.
+**Phase 1 — Encodeur (auto-supervisé, sans labels d'opt)** — `train_jepa.py`
 
----
+JEPA latent classique transposé sur le graphe : on masque des nœuds et on prédit
+leur **représentation** (pas leurs mnémoniques). Encodeur en ligne `θ` + encodeur
+cible `ξ = EMA(θ)`, stop-gradient, et garde-fou anti-collapse VICReg
+(variance + covariance). Tous les blocs (O0…O3) sont dans le même tas, vus comme
+des programmes bruts.
 
-### Local / generic (start here)
+**Phase 2 — Prédicteur (sur l'encodeur gelé)** — `train_predictor.py`
 
-We use [uv](https://docs.astral.sh/uv/guides/projects/) for package management.
+C'est ici, et seulement ici, qu'on utilise les paires O0↔O3 de la même fonction :
 
-```bash
-# Install dependencies
-uv sync
-# Option 1: Activate virtual environment
-source .venv/bin/activate
-python -m examples.image_jepa.main
-# Option 2: Run directly with uv
-uv run python -m examples.image_jepa.main
 ```
-If you need conda-specific packages, you can use **Conda + uv**
-
-```bash
-# Create conda environment with Python 3.12
-conda create -n eb_jepa python=3.12 -y
-conda activate eb_jepa
-# Install package in editable mode with dev dependencies (pytest, black, isort, autoflake)
-uv pip install -e . --group dev
+état   = un programme          → g = encodeur_gelé(programme)
+action = un niveau d'opt cible  → embedding d'action appris
+modèle :  g(forme_i) + action(niveau_j) → prédiction de g(forme_j)   (résidu = l'« édit »)
 ```
 
-Add these to your `~/.bashrc` for persistent configuration.
-
-```bash
-# Where datasets are stored / looked up
-export EBJEPA_DSETS=/path/to/eb_jepa/datasets
-# Optional: Directory for checkpoints and logs
-export EBJEPA_CKPTS=/path/to/checkpoints
-```
+**Métrique honnête : battre l'identité.** Comme le sens est ~invariant, `g(O0)`
+est déjà proche de `g(O3)` ; le prédicteur n'a de valeur que s'il retrouve le vrai
+`g(O3)` *mieux* que le no-op `g(O0)`. On rapporte la retrieval@1 des deux + la
+similarité cosinus.
 
-Verify the install with `uv run pytest tests/`.
+---
 
-## 🏋️ Training
+## 🚀 Démarrage rapide
 
-### Quick Start
+On utilise [uv](https://docs.astral.sh/uv/) pour la gestion des dépendances.
+Le pipeline asm a besoin de **clang** sur le `PATH`.
 
 ```bash
-# Local training
-python -m examples.{image_jepa,video_jepa,ac_video_jepa}.main
-```
-> Our default configs are tuned for H100 GPUs. With older GPUs (e.g., A100, V100), you may need to reduce batch size to fit in memory.
-
-### 📂 Folder Structure
+# 1. Dépendances
+uv sync
+source .venv/bin/activate
 
-All experiments use a unified folder structure:
+# 2. Générer un corpus C, le compiler et construire la base de graphes
+python -m examples.asm_superopt.gen_corpus --n 4000   # → fichiers .c
+python -m eb_jepa.asm.corpus                          # → data/corpus.jsonl
 
-```
-checkpoints/
-└── {example_name}/
-    ├── dev_2026-01-16_00-10/                 # Single/local runs (dev_ prefix)
-    │   └── {exp_name}_seed1/
-    │
-    ├── sweep_2026-01-16_00-10/         # Auto-named 3-seed sweep
-    │   ├── {exp_name}_seed1/
-    │   ├── {exp_name}_seed1000/
-    │   └── {exp_name}_seed10000/
-    │
-    └── sweep_my_experiment/            # Custom-named sweep
-        └── ...
+# 3. Pipeline complet (encodeur → prédicteur → figures de diagnostic)
+bash examples/asm_superopt/run_all.sh
 ```
 
-`{exp_name}` encodes key hyperparameters to avoid folder collisions, e.g.:
-- **image_jepa**: `resnet_vicreg_proj_bs256_ep300_ph2048_po2048_std1.0_cov80.0`
-- **video_jepa**: `resnet_bs64_lr0.001_std10.0_cov100.0`
-- **ac_video_jepa**: `impala_cov8_std16_simt12_idm1`
+> `run_all.sh` tourne sur **CPU** par défaut : pour ce petit GNN, le fallback CPU
+> des ops de scatter sur MPS rend le GPU local ~250× plus lent. Le vrai chemin
+> CUDA est le cluster SLURM (Vivatech / HTW).
 
-<details>
-<summary><span style="font-size: 1.17em; font-weight: bold;">🖥️ SLURM Launcher (optional)</span></summary>
+### Scripts du dossier `examples/asm_superopt/`
 
-| Command | Description |
-|---------|-------------|
-| `--example {name}` | Choose: `image_jepa`, `video_jepa`, `ac_video_jepa`, `maze`, `fintime`, `ltsf`, `eeg`, `audio`, `pointcloud`, `gray_scott`, `intuitive_physics`, `factors_of_variation` |
-| `--fname {path}` | Run the sweep specified in the config at `{path}` |
-| `--single` | Launch single job (dev mode) |
-| `--sweep {name}` | Custom sweep name |
-| `--array-parallelism {N}` | Limits the maximum number of concurrent jobs to `N` |
-| `--full-sweep` | Full hyperparameter sweep from config |
-| `--use-wandb-sweep` | Enable wandb sweep UI |
+| Script | Rôle |
+|---|---|
+| `gen_corpus.py` | génère un corpus C de fonctions feuilles straight-line |
+| `train_jepa.py` | Phase 1 — encodeur JEPA latent (masking, sans labels d'opt) |
+| `train_mask.py` | baseline générative (prédit les mnémoniques masqués) |
+| `train_predictor.py` | Phase 2 — prédicteur world-model sur encodeur gelé |
+| `analyze_opt_proximity.py` | diagnostic : O0/O1/O2/O3 d'une même fonction sont-ils proches ? |
+| `compare_idempotence.py` | le prédicteur est-il idempotent (O3 → O3 fixe) ? |
+| `study_o3_input.py` | comportement du prédicteur selon le niveau d'entrée |
+| `bridge_specificity.py` | le « pont » O0→O3 est-il spécifique à la bonne fonction ? |
 
-```bash
-# 3 seeds with wandb averaging (recommended)
-python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml
-
-# Custom sweep name
-python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --sweep my_experiment
-
-# Single job
-python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --single
+---
 
-# Full hyperparameter sweep
-python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --full-sweep
+## 📁 Structure
 
-# With wandb sweep UI for hyperparameter analysis
-python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --use-wandb-sweep
 ```
-
-Replace `image_jepa` with `ac_video_jepa`, `video_jepa`, or `maze` for other examples.
-
-**Full Sweep Configuration:** The `--full-sweep` flag reads the `sweep.param_grid` section from the example's YAML config file (e.g., `examples/image_jepa/cfgs/default.yaml`). Without this flag, only a 3-seed sweep is launched. To customize sweep parameters, edit the `sweep` section in the config:
-
-```yaml
-# Example: examples/image_jepa/cfgs/default.yaml
-sweep:
-  param_grid:
-    loss.cov_coeff: [0.1, 1.0, 10.0, 100.0]
-    loss.std_coeff: [1.0, 10.0]
-    meta.seed: [1, 1000, 10000]
+eb_jepa/asm/                 cœur du projet (lib)
+├── dataflow.py              assembleur x86-64 → graphe de flot de données
+├── encoder.py               encodeur GNN factorisé (meaning / temporality)
+├── corpus.py                .c → clang → blocs → base de graphes JSONL
+└── synth.py                 blocs synthétiques + réécritures préservant le sens
+
+examples/asm_superopt/       expériences, entraînement, diagnostics
+├── seed_corpus/  test_corpus/   sources C
+├── data/                    corpus JSONL, checkpoints (.pt), figures
+└── run_all.sh               orchestrateur retrain + ré-évaluation
+
+eb_jepa/                     librairie JEPA upstream (jepa, losses, planning, …)
+tests/                       test_asm_dataflow.py · test_asm_encoder.py · test_asm_corpus.py
+docs/                        slides, schémas, barème
 ```
 
-### Wandb Seed Averaging
-
-Runs with the same hyperparameters but different seeds share the same wandb run name, enabling automatic averaging:
-
-1. Go to wandb web UI → Runs table
-2. Click **"Group by"** → select **"Name"**
-   → Groups runs with identical hyperparameters (different seeds) together
-
-To filter runs from a specific sweep:
-3. Click **"Filter"** → **"Group"** → select your sweep name
-
-For detailed wandb sweep analysis (parallel coordinates, hyperparameter importance):
-1. Use `--use-wandb-sweep` flag when launching
-2. Go to wandb web UI → left pane → **"Sweeps"** → click your sweep name
-
-**SLURM Configuration:** SLURM parameters default to the HTW cluster and are read from `EBJEPA_SLURM_*` env vars (set by `env.sh`, which also auto-detects your account/QOS per user). Override per launch with the CLI flags `--partition`/`--account`/`--cpus-per-task`/`--time-min`/`--gpus-per-node`, or export the matching `EBJEPA_SLURM_*` var. The `SLURM_DEFAULTS` dictionary at the top of `examples/launch_sbatch.py` holds the fallbacks.
-
-</details>
-
-## 🧪 Running test cases
+---
 
-Libraries added to eb_jepa [must have their own test cases](/tests/). To run the tests:
+## 🧪 Tests
 
 ```bash
-# With uv sync installation
 uv run pytest tests/
-# With conda + uv installation (no .venv created)
-pytest tests/
 ```
 
-## 👩‍💻 Development
+Les tests propres au projet : `test_asm_dataflow.py`, `test_asm_encoder.py`,
+`test_asm_corpus.py`.
 
-Before contributing, please format your code with the following tools:
+## 👩‍💻 Développement
+
+Avant de contribuer, formatez le code :
 
 ```bash
-# Remove unused imports
 autoflake --remove-all-unused-imports -r --in-place .
-# Sort imports
 python -m isort eb_jepa examples tests
-# Format code
 python -m black eb_jepa examples tests
 ```
 
-## 📚 Citing EB-JEPA
+---
+
+## 📚 Crédits & upstream
 
-If you find this repository useful, please consider giving a ⭐ and citing:
+Ce dépôt est construit sur **EB-JEPA** (*Energy-Based Joint-Embedding Predictive
+Architectures*), la librairie open-source de **Meta AI Research (FAIR)** :
+Basile Terver, Randall Balestriero, Megi Dervishi, David Fan, Quentin Garrido,
+Tushar Nagarajan, Koustuv Sinha, Wancong Zhang, Mike Rabbat, Yann LeCun, Amir Bar.
 
 ```bibtex
 @misc{terver2026lightweightlibraryenergybasedjointembedding,
@@ -250,6 +195,6 @@ If you find this repository useful, please consider giving a ⭐ and citing:
 }
 ```
 
-## 📄 License
+## 📄 Licence
 
-EB-JEPA is Apache licensed. See [LICENSE](LICENSE.md).
+Sous licence Apache 2.0 (héritée d'EB-JEPA). Voir [LICENSE](LICENSE.md).