⚙️ JEPA-ASML

JEPA pour la superoptimisation d'assembleur

Apprendre une représentation du sens d'un programme assembleur — invariante à la forme produite par le compilateur — et un prédicteur world-model qui transforme la représentation d'un code lent (O0) en celle de son équivalent optimisé (O3).

---

🎯 Le projet en une page

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.

L'idée de ce projet est d'appliquer les Joint-Embedding Predictive Architectures (JEPA) à ce problème, avec une représentation factorisée :

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)

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.

Ce dépôt est un fork de la librairie EB-JEPA 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.

---

🧩 Le pipeline

  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)

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.

---

🏋️ Entraînement en deux phases

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.

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.

Phase 2 — Prédicteur (sur l'encodeur gelé) — train_predictor.py

C'est ici, et seulement ici, qu'on utilise les paires O0↔O3 de la même fonction :

é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 »)

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.

---

🚀 Démarrage rapide

On utilise uv pour la gestion des dépendances. Le pipeline asm a besoin de clang sur le PATH.

# 1. Dépendances
uv sync
source .venv/bin/activate

# 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

# 3. Pipeline complet (encodeur → prédicteur → figures de diagnostic)
bash examples/asm_superopt/run_all.sh

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).

Scripts du dossier examples/asm_superopt/

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 ?

---

📁 Structure

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

---

🧪 Tests

uv run pytest tests/

Les tests propres au projet : test_asm_dataflow.py, test_asm_encoder.py, test_asm_corpus.py.

👩‍💻 Développement

Avant de contribuer, formatez le code :

autoflake --remove-all-unused-imports -r --in-place .
python -m isort eb_jepa examples tests
python -m black eb_jepa examples tests

---

📚 Crédits & upstream

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.

@misc{terver2026lightweightlibraryenergybasedjointembedding,
      title={A Lightweight Library for Energy-Based Joint-Embedding Predictive Architectures},
      author={Basile Terver and Randall Balestriero and Megi Dervishi and David Fan and Quentin Garrido and Tushar Nagarajan and Koustuv Sinha and Wancong Zhang and Mike Rabbat and Yann LeCun and Amir Bar},
      year={2026},
      eprint={2602.03604},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2602.03604},
}

📄 Licence

Sous licence Apache 2.0 (héritée d'EB-JEPA). Voir LICENSE.