NetTAG: A Multimodal RTL-and-Layout-Aligned Netlist Foundation Model via Text-Attributed Graph

NetTAG represents a synthesized gate-level netlist as a text-attributed graph (TAG) in which every node carries a Boolean-expression embedding produced by a fine-tuned LLM (ExprLLM), and then learns a multimodal foundation model (TAGFormer) that aligns the netlist TAG with its corresponding RTL text and post-place-and-route layout graph. The resulting netlist embeddings transfer to a range of downstream EDA tasks (functional block classification, timing prediction, GNN-RE reverse engineering, etc.) without re-running expensive flows.

This document is a step-by-step walkthrough of how to use the repository — from raw .v netlists to a fine-tuned downstream model. For a higher-level summary of the architecture, see CLAUDE.md.

---

Table of contents

1. Pipeline overview
2. Installation
3. Repository layout
4. Stage 0 — Data collection
5. Stage 1 — Preprocess (cone extraction per modality)
6. Stage 2 — Train ExprLLM (the node text encoder)
7. Stage 3 — Apply ExprLLM to build the TAG dataset
8. Stage 4 — Pretrain TAGFormer (multimodal alignment)
9. Stage 5 — Inference and downstream fine-tuning
10. Example: GNN-RE benchmark end-to-end
11. Tips, gotchas, and FAQ

---

1. Pipeline overview

                  ┌──────────────────┐
 raw netlists ─►  │  preprocess /    │ ─►  per-design graphs +
 (.v, PT rpts)    │  net2graph_*     │     per-cone subgraphs
                  └──────────────────┘                  │
                                                        ▼
                  ┌──────────────────┐         expression strings
                  │  expr_aug /      │ ─►      per node (TAG text)
                  │  aug.py          │
                  └──────────────────┘                  │
                                                        ▼
                  ┌──────────────────┐
                  │  model/exprllm   │ ─►  LoRA-tuned LLM that
                  │  (LLM2Vec stack) │     embeds each expression
                  └──────────────────┘                  │
                                                        ▼
                  ┌──────────────────┐
                  │  subgraph2       │ ─►  PyG Data objects with
                  │  dataset_tag.py  │     per-node feature vectors
                  └──────────────────┘                  │
                                                        ▼
            ┌─────────────────────────────────────┐
            │  model/tagformer                    │ ─►  netlist
            │  (RTL_Fusion + Net_Encoder, GT/GNN, │     embeddings,
            │   xbert; aligns net / RTL / layout) │     downstream
            └─────────────────────────────────────┘     fine-tune

The four modalities the model sees are:

---

2. Installation

NetTAG assumes a CUDA-capable Linux machine (multiple GPUs strongly recommended for ExprLLM training). The code has been used with Python 3.10, CUDA 12.x, and PyTorch 2.x.

# 1. Clone the repo
git clone <this-repo-url> NetTAG
cd NetTAG

# 2. (Recommended) create an env
conda create -n nettag python=3.10 -y
conda activate nettag

# 3. Install ExprLLM / LLM2Vec wrapper (editable)
cd model/exprllm
pip install -e .
pip install flash-attn --no-build-isolation
cd ../..

# 4. Other dependencies used across the pipeline
pip install \
  torch torchvision torchaudio \
  torch_geometric dgl \
  transformers accelerate peft \
  sentence-transformers \
  pyverilog pysmt \
  networkx numpy scipy scikit-learn \
  xgboost \
  ruamel.yaml tqdm tensorboard

The model/exprllm/llm2vec/ package is a fork of McGill-NLP/llm2vec; see its own model/exprllm/README.md for the upstream usage notes.

External tools. A pre-existing synthesis + STA flow is required to generate the inputs in data_collect/data_pt, data_pt_pos, and the PT timing reports. NetTAG itself does not run synthesis; it consumes the artifacts.

---

3. Repository layout

NetTAG/
├── data_collect/                    # raw inputs to the pipeline
│   ├── data_pt → ../../net_tag/...  # netlists + PT reports (symlink)
│   ├── data_pt_pos → ...            # post-PnR variant
│   ├── data_gnnre/                  # GNN-RE benchmark netlists
│   └── data_js_*/                   # per-dataset design lists (JSON)
│
├── preprocess/
│   ├── net2graph_ori/               # baseline: netlist → graph + cones
│   ├── net2graph_pos/               # post-place-and-route variant
│   ├── net2graph_layout/            # physical-aware layout graph
│   ├── net2graph_gnnre/             # GNN-RE benchmark variant
│   ├── net2graph_gnnre_aig/         # GNN-RE AIG variant
│   └── rtl2embed/                   # RTL text → embedding (NV-Embed-v1)
│
├── model/
│   ├── exprllm/                     # ExprLLM = LLM2Vec fork
│   │   ├── llm2vec/                 # upstream library code
│   │   ├── exprllm/                 # NetTAG-specific train/infer scripts
│   │   ├── train_configs/exprllm/   # pretraining configs (4-stage curriculum)
│   │   └── infer_configs/exprllm/   # inference / node-dict update configs
│   └── tagformer/                   # multimodal fusion model
│       ├── models/                  # Net_Encoder, RTL_Fusion, xbert, gt, gnn
│       ├── configs/                 # YAML + per-submodel JSONs
│       ├── pretrain_net*.py         # pretraining entry points
│       ├── finetune_task{1..4}.py   # per-task fine-tuning heads
│       ├── net2vec_task*.py         # inference (embeddings only)
│       └── dataset_proc/            # dataset loaders + symlink instructions
│
├── dataset/                         # produced datasets (and symlinks)
│   ├── dataset_pretrain_align/      # train/valid split scripts
│   ├── dataset_finetune/            # downstream-task split scripts
│   ├── design/, expr/               # symlinks into ../../net_tag/dataset/
│   └── graph/{gnnre,layout}/        # final PyG-ready graph dataset
│
├── README.md                        # ← this file
└── CLAUDE.md                        # high-level notes for Claude Code

---

4. Stage 0 — Data collection

data_collect/ holds the raw artifacts you produce outside this repo:

• data_collect/data_pt/<design>/ — per-design synthesized netlist (*.v) and PrimeTime reports (timing / power / capacitance).
• data_collect/data_pt_pos/<design>/ — same structure for the post-place-and-route version.
• data_collect/data_gnnre/ — gate-level netlists for the GNN-RE benchmark (one .v per design).
• data_collect/data_js_<dataset>/design_list.json — flat JSON array of design names that every run_parallel.py reads. For benchmarks with predefined splits there are also train_list.json, val_list.json, test_list.json (e.g. data_js_gnnre/).

If the symlinks in data_collect/data_pt* point at ../../net_tag/data_collect/..., re-create them to wherever your synthesized data actually lives.

---

5. Stage 1 — Preprocess (cone extraction per modality)

Each preprocess/net2graph_* directory contains the same skeleton:

net2graph_ori/
├── pyverilog/                    # vendored Pyverilog (used by AST_analyzer)
├── parse_split_template/
│   └── net2graph_template.py     # template with DESIGN_NAME_HERE placeholders
├── AST_analyzer.py               # walks the Verilog AST → DAG
├── DG.py                         # graph utility (NetworkX-based)
├── run_parallel.py               # multiprocessing driver
├── expr_aug/                     # Boolean-expression augmentation (pySMT)
├── subgraph2dataset_tag.py       # cone subgraph → PyG Data object
├── saved_analyzer/               # ← outputs of step 1 (per design)
├── saved_graph_split/            # ← outputs of step 1 (per cone)
└── saved_expr/                   # ← outputs of step 2

5.1 Netlist → graph + per-register cone

run_parallel.py reads data_collect/data_js_<dataset>/design_list.json, then for each design:

1. Copies parse_split_template/ to a temp dir.
2. Substitutes DESIGN_NAME_HERE inside net2graph_template.py to produce net2graph.py.
3. Runs the resulting net2graph.py (which walks the Verilog AST, builds a NetworkX DAG, and slices the DAG into per-register cone subgraphs).
4. Deletes the temp dir.

cd preprocess/net2graph_ori
python3 run_parallel.py

Outputs

• saved_analyzer/<design>/: pickled AST/graph for the whole design.
• saved_graph_split/<design>/<cone>.pkl: per-cone subgraph + node_dict mapping node IDs to their symbolic Boolean expression.

⚠️ Edit parse_split_template/net2graph_template.py if you need to change preprocessing — not the generated net2graph.py, which is regenerated and deleted on every run.

⚠️ run_parallel.py uses multiprocessing.Pool(50). Lower the pool size on machines with fewer CPUs / less RAM.

5.2 Build Boolean-expression augmentation pairs

For ExprLLM contrastive training we need pairs of equivalent Boolean expressions (a positive) and non-equivalent ones (a negative). expr_aug/aug.py walks saved_graph_split/ and uses pySMT to apply logical-equivalence-preserving rewrites:

cd preprocess/net2graph_ori/expr_aug
python3 aug.py

Outputs

• ../saved_expr/<design>/<cone>.pkl: pairs/triplets of expressions consumed by the ExprLLM pretraining config (dataset_file_path in train_configs/exprllm/*.json).

5.3 Layout-graph variant

preprocess/net2graph_layout/ reuses the same Step 1, but subgraph2dataset.py packs each cone into a graph whose nodes carry physical attributes parsed from PT reports rather than expression embeddings. Run it after you also have saved_graph_split/ for the layout flow.

5.4 RTL text embeddings

Each cone's RTL source code is embedded with nvidia/NV-Embed-v1:

cd preprocess/rtl2embed/scr
python3 vlg2vec_nv.py

The script expects per-design *.json files listing the cone endpoints (default path inside the script — edit folder_dir and save_dir to match your layout).

---

6. Stage 2 — Train ExprLLM (the node text encoder)

ExprLLM is a fork of LLM2Vec: a decoder-only LLM (Sheared-LLaMA-1.3B by default) is converted into a bidirectional text encoder via masked-next-token-prediction + contrastive learning on the expression pairs from §5.2. NetTAG uses a 4-stage curriculum controlled by the four configs in model/exprllm/train_configs/exprllm/Sheared-Llama_{1..4}.json.

cd model/exprllm/exprllm
python3 run_exprllm_pretrain.py ../train_configs/exprllm/Sheared-Llama_1.json
python3 run_exprllm_pretrain.py ../train_configs/exprllm/Sheared-Llama_2.json
python3 run_exprllm_pretrain.py ../train_configs/exprllm/Sheared-Llama_3.json
python3 run_exprllm_pretrain.py ../train_configs/exprllm/Sheared-Llama_4.json

Each config sets:

• model_name_or_path: base LLM (default princeton-nlp/Sheared-LLaMA-1.3B).
• peft_model_name_or_path: LoRA adapter from the previous stage (skipped in stage 1).
• dataset_file_path: a .pkl produced by expr_aug/aug.py.
• LoRA hyperparameters, attn_implementation: flash_attention_2, torch_dtype: bfloat16, max_seq_length: 8192.

⚠️ The training scripts hardcode CUDA_VISIBLE_DEVICES = "0,1,2,3,4,5" at the top of run_exprllm_pretrain.py. Change the literal there if you don't have 6 GPUs.

⚠️ The scripts also sys.path.append('/home/usr/NetTAG/model/exprllm'). Replace /home/usr/ with your install path on first run, otherwise the in-repo llm2vec will not be picked up.

Outputs land in output_dir from the config — by default model/exprllm/output/exprllm_1_5B/part{1..4}/.

---

7. Stage 3 — Apply ExprLLM to build the TAG dataset

With ExprLLM trained, we replace each node's expression string in saved_graph_split/ with its embedding vector. The driver script lives next to the training script:

cd model/exprllm/exprllm
python3 run_update_node_dict.py ../infer_configs/exprllm/Sheared-Llama.json

The inference config (infer_configs/exprllm/Sheared-Llama.json) points at the final training checkpoint via peft_model_name_or_path (e.g. output/exprllm_1_5B/part4/checkpoint-6879). Update that path before running.

Outputs

• save_node_dict_tag/<design>/<cone>.pkl: updated node_dict containing both the original expression and its feature vector.

Then materialize the PyG dataset:

cd preprocess/net2graph_ori
python3 subgraph2dataset_tag.py

This script (see subgraph2dataset_tag.run_one_subgraph) injects a synthetic [CLS] node connected to every other node, builds a torch_geometric.data.Data per cone, and saves the result under dataset/tag/ori/ (or pos/, layout/, gnnre/ from the corresponding variant).

For the layout modality, run preprocess/net2graph_layout/subgraph2dataset.py (note: no _tag suffix — the layout graph does not use ExprLLM features). For the GNN-RE benchmark, use net2graph_gnnre/subgraph2dataset_tag.py (and net2graph_gnnre_aig/ for AIG-form netlists).

---

8. Stage 4 — Pretrain TAGFormer (multimodal alignment)

TAGFormer is the cross-modal fusion model (model/tagformer/). Architecturally:

• models/model_net.py::Net_Encoder — node-level encoder over the TAG (a Graph Transformer from models/gt.py + GNN from models/gnn.py).
• models/xbert.py — cross-attention BERT used to fuse netlist embeddings with RTL text embeddings.
• models/model_pretrain.py::RTL_Fusion — the joint pretraining wrapper that combines Net_Encoder, RTL embeddings, and the layout encoder under a contrastive + masked-modeling objective.

8.1 Prepare the dataset symlinks

dataset_proc/load_dataset.py expects a fixed set of file names under model/tagformer/dataset_proc/. The exact ln -s commands are in dataset_proc/README.md — they point at dataset/{net,rtl}_data/data_bench/dataset_{train,valid,test,sft}{,1..4}_{ori,pos,neg}.pkl.

To produce those .pkl files first, run the split scripts:

cd dataset/dataset_pretrain_align
python3 split_train_valid_align.py       # full pretrain/align split
python3 split_train_valid_align_part.py  # chunked variant (train1..train4)
python3 split_train_valid_net.py         # netlist-only variant

For downstream fine-tuning, generate per-design test splits with dataset/dataset_finetune/split_one_design*.py.

8.2 Configure

configs/Pretrain.yaml controls high-level pretraining:

bert_model_name: 'bert-base-uncased'
bert_config: 'configs/config_bert.json'
gt_config:   'configs/config_gt.json'
gnn_config:  'configs/config_gnn.json'

text_width: 4096       # NV-Embed-v1 embedding width
embed_dim:  768
batch_size: 2
temp:       0.07       # InfoNCE temperature
queue_size: 65536      # MoCo-style negative queue
momentum:   0.995

optimizer: { lr: 1e-4, eps: 1e-8, weight_decay: 0 }
schedular: { epochs: 150, lr_end: 1e-5, warmup_lr: 1e-5,
             warmup_updates: 20, total_updates: 256 }

Sub-encoder configs (config_bert.json, config_gt.json, config_gnn.json) tweak depth/heads/hidden size of each backbone independently.

8.3 Run pretraining

There are three entry points; pick the one that matches the modalities you have prepared:

cd model/tagformer
accelerate launch pretrain_net_align.py --config configs/Pretrain.yaml
# or plain `python3 pretrain_net_align.py --config configs/Pretrain.yaml`

pretrain_net_align.py uses 🤗 Accelerate; it forces CUDA_VISIBLE_DEVICES = "0,1" at the top of the file — edit if needed. TensorBoard logs are written via torch.utils.tensorboard.SummaryWriter.

Checkpoints land under output/<date>/ where date is set near the top of each script (e.g. date = 'pretrain_net_align_7B_1024').

---

9. Stage 5 — Inference and downstream fine-tuning

9.1 Extract netlist embeddings only

For each task there is a net2vec_task<i>.py script that loads a pretrained Net_Encoder checkpoint and writes one embedding per cone:

cd model/tagformer
python3 net2vec_task1.py     # task 1
python3 net2vec_task3.py     # task 3
python3 net2vec_task4.py     # task 4 (functional block identification)
python3 net2vec_task4_aig.py # task 4 on AIG-form netlists

Use these embeddings as features for any external classifier / regressor.

9.2 End-to-end fine-tune

finetune_task<i>.py attaches a task-specific head on top of Net_Encoder (typically an XGBRegressor / XGBClassifier or an MLPRegressor / MLPClassifier from sklearn — see finetune_task1.py:18 for the imports) and trains end-to-end:

python3 finetune_task1.py    # Stage-1 regression task
python3 finetune_task2.py    # Stage-2
python3 finetune_task3.py    # Stage-3
python3 finetune_task4.py    # GNN-RE functional block classification
python3 finetune_task4_aig.py

Metrics (utils/eval.py::regression_metrics, classify_metrics) and per-design test loaders (dataset_proc.load_test_dataset_finetune_one_design) are wired into each script.

---

10. Example: GNN-RE benchmark end-to-end

A complete reproduction of the GNN-RE downstream experiment:

# 0. Place GNN-RE netlists in data_collect/data_gnnre/ and update
#    data_collect/data_js_gnnre/{design_list,train_list,val_list,test_list}.json

# 1. Cone extraction
cd preprocess/net2graph_gnnre
python3 run_parallel.py

# 2. Build ExprLLM training pairs (only needed once across datasets)
cd ../net2graph_ori/expr_aug && python3 aug.py

# 3. Train ExprLLM (4 stages)
cd ../../../model/exprllm/exprllm
for i in 1 2 3 4; do
    python3 run_exprllm_pretrain.py ../train_configs/exprllm/Sheared-Llama_${i}.json
done

# 4. Populate node features on GNN-RE cones
python3 run_update_node_dict_gnnre.py ../infer_configs/exprllm/Sheared-Llama.json

# 5. Materialize PyG dataset
cd ../../../preprocess/net2graph_gnnre
python3 subgraph2dataset_tag.py

# 6. Symlink pkl files into model/tagformer/dataset_proc/
#    (see model/tagformer/dataset_proc/README.md)

# 7. Pretrain TAGFormer (or load a public checkpoint)
cd ../../model/tagformer
accelerate launch pretrain_net_align.py --config configs/Pretrain.yaml

# 8. Fine-tune the task-4 head
python3 finetune_task4.py

---

11. Tips, gotchas, and FAQ

Hardcoded paths and GPU lists

Almost every Python entry point sets os.environ["CUDA_VISIBLE_DEVICES"] at the top of the file (it has to happen before import torch), and several add absolute sys.path entries like /home/usr/NetTAG/.... Always check the top of a script before running it on a new machine:

grep -n -E 'CUDA_VISIBLE_DEVICES|sys.path.append|/home/' <script>.py

Multiprocessing pool sizes

preprocess/*/run_parallel.py defaults to Pool(50). Each worker spawns Pyverilog/pySMT and parses an entire design — easily 1–2 GB per worker on large designs. Lower to Pool(8) or Pool(16) if your machine has <128 GB RAM.

Editing the graph parser

parse_split_template/net2graph_template.py is the template; net2graph.py next to it (when present) is generated and will be deleted on the next run. Always edit the template.

"Where does X come from?"

Reproducibility

There is no pytest suite or CI in this repo. Validation is "did stage N produce the expected output directory?" The fastest smoke test is to run run_parallel.py over a 2-design subset (edit the design list temporarily) and then subgraph2dataset_tag.py on the resulting cones — if both finish without exceptions, the toolchain is wired correctly.

---

License and acknowledgements

model/exprllm/ is a fork of McGill-NLP/llm2vec (MIT). Pyverilog vendored under each preprocess/net2graph_*/pyverilog/ retains its original Apache 2.0 license. The rest of NetTAG-specific code is released under the terms in this repository.