NetTAG

KO | EN

gitlite — search

#python #hacktoberfest #react #typescript #go #golang #machine-learning #rust #linux #nodejs #java #windows #pytorch #ios

(★ 25)

NetTAG: A Multimodal RTL-and-Layout-Aligned Netlist Foundation Model via Text-Attributed Graph (DAC'25)

NetTAG 최신버젼 다운로드

최종 버전 다운로드 (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# 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](#1-pipeline-overview)
2. [Installation](#2-installation)
3. [Repository layout](#3-repository-layout)
4. [Stage 0 — Data collection](#4-stage-0--data-collection)
5. [Stage 1 — Preprocess (cone extraction per modality)](#5-stage-1--preprocess-cone-extraction-per-modality)
6. [Stage 2 — Train ExprLLM (the node text encoder)](#6-stage-2--train-exprllm-the-node-text-encoder)
7. [Stage 3 — Apply ExprLLM to build the TAG dataset](#7-stage-3--apply-exprllm-to-build-the-tag-dataset)
8. [Stage 4 — Pretrain TAGFormer (multimodal alignment)](#8-stage-4--pretrain-tagformer-multimodal-alignment)
9. [Stage 5 — Inference and downstream fine-tuning](#9-stage-5--inference-and-downstream-fine-tuning)
10. [Example: GNN-RE benchmark end-to-end](#10-example-gnn-re-benchmark-end-to-end)
11. [Tips, gotchas, and FAQ](#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:

| Modality      | Producer script                              | What it captures                                                    |
|---------------|----------------------------------------------|----------------------------------------------------------------------|
| Netlist TAG   | `preprocess/net2graph_ori`                   | Gate-level graph + per-node Boolean expression                       |
| Layout graph  | `preprocess/net2graph_layout`                | Same graph topology + physical attributes from PT reports            |
| RTL text      | `preprocess/rtl2embed/scr/vlg2vec_nv.py`     | NV-Embed-v1 embedding of the RTL source code of each cone            |
| Positional    | `preprocess/net2graph_pos`                   | Position-aware variant of the netlist graph (post-place-and-route)   |

---

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

```bash
# 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](https://github.com/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.

```bash
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:

```bash
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`](https://huggingface.co/nvidia/NV-Embed-v1):

```bash
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](https://arxiv.org/abs/2404.05961): 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`.

```bash
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:

```bash
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:

```bash
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:

```bash
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:

```yaml
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:

| Script                  | What it trains                                                     |
|-------------------------|--------------------------------------------------------------------|
| `pretrain_net.py`       | Netlist-only pretraining (graph contrastive + masked node modeling)|
| `pretrain_net_align.py` | Full netlist ↔ RTL ↔ layout alignment (the headline NetTAG model)  |
| `pretrain_net_all.py`   | Stage-wise combined training                                       |

```bash
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:

```bash
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:

```bash
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:

```bash
# 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:

```bash
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?"

| Output file/dir                                | Produced by                                           |
|------------------------------------------------|--------------------------------------------------------|
| `saved_analyzer/`                              | `preprocess/net2graph_*/run_parallel.py`               |
| `saved_graph_split/`                           | `preprocess/net2graph_*/run_parallel.py`               |
| `saved_expr/`                                  | `preprocess/net2graph_ori/expr_aug/aug.py`             |
| `model/exprllm/output/exprllm_*/`              | `model/exprllm/exprllm/run_exprllm_pretrain.py`        |
| `save_node_dict_tag/`                          | `model/exprllm/exprllm/run_update_node_dict*.py`       |
| `dataset/tag/{ori,pos,gnnre,...}/`             | `preprocess/net2graph_*/subgraph2dataset_tag.py`       |
| `dataset/graph/layout/`                        | `preprocess/net2graph_layout/subgraph2dataset.py`      |
| `preprocess/rtl2embed/embeds/`                 | `preprocess/rtl2embed/scr/vlg2vec_nv.py`               |
| `model/tagformer/output/<date>/`               | `model/tagformer/pretrain_net*.py`                     |

### 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](https://github.com/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.

NetTAG

NetTAG 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

NetTAG

NetTAG 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

Download README