{"slug": "a-library-to-make-gnn-work-on-binary-code", "title": "A library to make GNN work on binary code", "summary": "Quarkslab has released pcode_graph, a Python library for building semantic graphs from binary code, and demonstrated its use in training a neural network to detect function similarities across architectures and compilers. The library, available on GitHub and via pip, was used in research presented at ESANN 2026 and is intended for tasks such as vulnerability detection and binary diffing.", "body_md": "Author\n\n[Samuel Hangouët](./author/samuel-hangouet.html)\n\nCategory\n\n[Program Analysis](./category/program-analysis.html)\n\nTags\n\n[machine learning](./tag/machine-learning.html),\n\n[function similarity](./tag/function-similarity.html),\n\n[data analysis](./tag/data-analysis.html),\n\n[binary analysis](./tag/binary-analysis.html),\n\n[GNN](./tag/gnn.html),\n\n[tool](./tag/tool.html),\n\n[graph neural networks](./tag/graph-neural-networks.html),\n\n[2026](./tag/2026.html)\n\n`pcode_graph`\n\nis a Python library, published by Quarkslab, suitable to build semantic graphs from binary code. We present how to use it to detect function similarities in binaries.\n\n## Context\n\nWe introduce here the Python library `pcode_graph`\n\n, a tool developed at Quarkslab to abstract the semantics of binary code. It provides an API to build and visualize Control & Data flow Graphs (CDG) from a function, a basic-block or any arbitrary piece of code.\n\nAs soon as you start looking into the automated analysis of binaries, you quickly realize that many use cases require extracting a semantic representation of a piece of code. For example, such a representation could be used to:\n\n- Identify the type(s) of obfuscation applied;\n- Build a database of gadgets or automatically chain them to do ROP;\n- Find changes between two successive versions of the same binary;\n- Look for a function in a database of binaries;\n- Look for vulnerabilities;\n- Deobfuscate a piece of binary.\n\nThis last use case was the subject of a [research paper published at ESANN 2026](https://github.com/quarkslab/conf-presentations/blob/master/Confs/ESANN2026/ES2026-155.pdf) using the `pcode_graph`\n\nlibrary, but it was not open-sourced at the time of publication.\n\nTo present the library, here we are going to focus on the detection of similarities between functions. More precisely, we will train a neural network to recognize a function, independently of the architecture, the compiler and the options used to compile it. For that, we will use the [Cisco-Talos dataset](https://www.usenix.org/system/files/sec22-marcelli.pdf).\n\nThe library is available from Quarkslab's `pcode_graph`\n\n[repository on GitHub](https://github.com/quarkslab/pcode_graph). It can also be installed directly with `pip`\n\n:\n\n``` bash\n(venv) $ pip install pcode_graph\n```\n\n## Extracting a semantic graph\n\nMany methods exist to teach a model to compare two binary functions.\n\nA basic approach is to extract statistical features such as the number of instructions, the number of basic blocks, the mnemonic frequencies, etc. This yields a table of features on which the learning is performed. This can be good enough to perform some basic tasks, like [malware classification](https://dl.acm.org/doi/pdf/10.1145/3576915.3616589).\n\n[jTrans](https://dl.acm.org/doi/pdf/10.1145/3533767.3534367) feeds the assembly (with a bit of preprocessing) into a language-processing model.\n\nFrom a production perspective, several tools were compared prior to creating Quarkslab's [Sighthouse tool](https://blog.quarkslab.com/sighthouse-automated-function-identification.html).\n\nHere we are going to extract a **semantic graph**, a representation of what the piece of code *does* disregarding *how* it does it, for example the specific CPU instructions. Note that this is not necessarily the most relevant approach: the choice depends entirely on your use case and your data.\n\nLet's take a very simple piece of code:\n\n```\nint do_it(int a, int b)\n{\n    if (a == b)\n        return a + b;\n    return 0;\n}\n```\n\nand compile it so as to get a small piece of assembly:\n\n```\nclang test.c -o test.o -Oz -c\n```\n\nOn x86_64 we get:\n\n```\nlea   ecx, [rsi + rdi*0x1]\nxor   eax, eax\ncmp   edi, esi\ncmovz eax, ecx\nret\n```\n\nNow let's try to compile it without optimization:\n\n```\npush rbp\nmov  rbp, rsp\nmov  dword ptr [rbp + -0x8], edi\nmov  dword ptr [rbp + -0xc], esi\nmov  eax, dword ptr [rbp + -0x8]\ncmp  eax, dword ptr [rbp + -0xc]\njnz  0x1d\nmov  eax, dword ptr [rbp + -0x8]\nadd  eax, dword ptr [rbp + -0xc]\nmov  dword ptr [rbp + -0x4], eax\njmp  0x24\nmov  dword ptr [rbp + -0x4], 0x0\nmov  eax, dword ptr [rbp + -0x4]\npop  rbp\nret\n```\n\nThese are two versions of the same function compiled with the same compiler for the same architecture, and yet:\n\n- The second version is three times longer.\n- The only mnemonic they have in common is CMP, and it works on different operands.\n- The second one has three basic blocks against a single one for the first.\n\nThese points highlight the limits of statistical feature extraction for function comparison.\n\nLet's now look at the optimized version. If we swap the two first instructions, the semantics remain unchanged:\n\n```\nxor   eax, eax\nlea   ecx, [rsi + rdi*0x1]\ncmp   edi, esi\ncmovz eax, ecx\nret\n```\n\nWhile the order of instructions can be of interest to detect the compiler used, for our use case we are to the contrary looking for a representation that is ideally identical for identical code semantics, otherwise our model would have to learn to ignore all the semantically equivalent permutations of a same function.\n\nTo do so, we are going to extract the data flow graph of the function.\n\nIn order to abstract away the architecture, we use the `pypcode`\n\nlibrary, a binding of `SLEIGH`\n\nwhich translates binary code into Ghidra's low-level internal representation. One benefit of P-Code is that it is limited to 63 distinct opcodes (at the *Raw* level, excluding `IMARK`\n\n). Compared to the thousands of x86_64 mnemonics, this limits the possibilities of expressing the same semantics in different ways. It is also a representation that makes building a data flow graph easier.\n\nWe can directly use the `pcode_graph`\n\nCLI to get the P-Code:\n\n```\ncdg pcode test.o > test.pcode\n```\n\nWe get the following file:\n\n``` php\nimark [0x0]\n$37632 = int_mult RDI, #0x1\n$38144 = int_add RSI, $37632\nECX = subpiece $38144, #0x0\nRCX = int_zext ECX\nimark [0x3]\nCF = #0x0\nOF = #0x0\nEAX = int_xor EAX, EAX\nRAX = int_zext EAX\nSF = int_sless EAX, #0x0\nZF = int_equal EAX, #0x0\n$361216 = int_and EAX, #0xff\n$361472 = popcount $361216\n$361728 = int_and $361472, #0x1\nPF = int_equal $361728, #0x0\nimark [0x5]\n$515328 = EDI\nCF = int_less $515328, ESI\nOF = int_sborrow $515328, ESI\n$515840 = int_sub $515328, ESI\nSF = int_sless $515840, #0x0\nZF = int_equal $515840, #0x0\n$361216 = int_and $515840, #0xff\n$361472 = popcount $361216\n$361728 = int_and $361472, #0x1\nPF = int_equal $361728, #0x0\nimark [0x7]\n$505344 = ECX\nRAX = int_zext EAX\n$505600 = bool_negate ZF\ncbranch [0xa], $505600\nEAX = $505344\nimark [0xa]\nRIP = load #0x6b970f0, RSP\nRSP = int_add RSP, #0x8\nreturn RIP\n```\n\nOk, that's fairly verbose... To break instructions down into basic operations, `pypcode`\n\nintroduces many temporary variables. But there is nothing to worry about: these variables will subsequently be ignored and will only make up edges of our data flow graph.\n\nBefore building the graph, we start by applying several analysis passes to the P-Code:\n\n- Instruction indexing (to resolve jumps);\n- Unreachable code detection;\n- Data flow analysis.\n\nThese results can be displayed as a table in markdown format:\n\n```\ncdg table test.o\n```\n\n| index | op | preds | succs | input defs | reachable | exit def |\n|---|---|---|---|---|---|---|\n| 0 | imark [0x0] | entry | 1 | x | ||\n| 1 | $37632 = int_mult RDI, #0x1 | 0 | 2 | RDI from entry | x | |\n| 2 | $38144 = int_add RSI, $37632 | 1 | 3 | RSI from entry, $37632 from 1 | x | |\n| 3 | ECX = subpiece $38144, #0x0 | 2 | 4 | $38144 from 2 | x | |\n| 4 | RCX = int_zext ECX | 3 | 5 | ECX from 3 | x | RCX |\n| 5 | imark [0x3] | 4 | 6 | x | ||\n| 6 | CF = #0x0 | 5 | 7 | x | ||\n| 7 | OF = #0x0 | 6 | 8 | x | ||\n| 8 | EAX = int_xor EAX, EAX | 7 | 9 | EAX from entry, EAX from entry | x | |\n| 9 | RAX = int_zext EAX | 8 | 10 | EAX from 8 | x | |\n| 10 | SF = int_sless EAX, #0x0 | 9 | 11 | EAX from 9 | x | |\n| 11 | ZF = int_equal EAX, #0x0 | 10 | 12 | EAX from 9 | x | |\n| 12 | $361216 = int_and EAX, #0xff | 11 | 13 | EAX from 9 | x | |\n| 13 | $361472 = popcount $361216 | 12 | 14 | $361216 from 12 | x | |\n| 14 | $361728 = int_and $361472, #0x1 | 13 | 15 | $361472 from 13 | x | |\n| 15 | PF = int_equal $361728, #0x0 | 14 | 16 | $361728 from 14 | x | |\n| 16 | imark [0x5] | 15 | 17 | x | ||\n| 17 | $515328 = EDI | 16 | 18 | EDI from entry | x | |\n| 18 | CF = int_less $515328, ESI | 17 | 19 | $515328 from 17, ESI from entry | x | CF |\n| 19 | OF = int_sborrow $515328, ESI | 18 | 20 | $515328 from 17, ESI from entry | x | OF |\n| 20 | $515840 = int_sub $515328, ESI | 19 | 21 | $515328 from 17, ESI from entry | x | |\n| 21 | SF = int_sless $515840, #0x0 | 20 | 22 | $515840 from 20 | x | SF |\n| 22 | ZF = int_equal $515840, #0x0 | 21 | 23 | $515840 from 20 | x | ZF |\n| 23 | $361216 = int_and $515840, #0xff | 22 | 24 | $515840 from 20 | x | |\n| 24 | $361472 = popcount $361216 | 23 | 25 | $361216 from 23 | x | |\n| 25 | $361728 = int_and $361472, #0x1 | 24 | 26 | $361472 from 24 | x | |\n| 26 | PF = int_equal $361728, #0x0 | 25 | 27 | $361728 from 25 | x | PF |\n| 27 | imark [0x7] | 26 | 28 | x | ||\n| 28 | $505344 = ECX | 27 | 29 | ECX from 4 | x | |\n| 29 | RAX = int_zext EAX | 28 | 30 | EAX from 9 | x | EAX |\n| 30 | $505600 = bool_negate ZF | 29 | 31 | ZF from 22 | x | |\n| 31 | cbranch [0xa], $505600 | 30 | 33, 32 | $505600 from 30 | x | |\n| 32 | EAX = $505344 | 31 | 33 | $505344 from 28 | x | EAX |\n| 33 | imark [0xa] | 32, 31 | 34 | x | ||\n| 34 | RIP = load #0x33fd1630, RSP | 33 | 35 | MEMORY from entry, RSP from entry | x | RIP |\n| 35 | RSP = int_add RSP, #0x8 | 34 | 36 | RSP from entry | x | RSP |\n| 36 | return RIP | 35 | elsewhere | RIP from 34 | x |\n\nSince `pcode_graph`\n\nwas initially designed to work on small chunks of code, the CFG is built directly at the level of the P-Code operations instead of extracting basic-blocks. Beware that this could change in the near future.\n\nFrom these elements the dataflow graph is easy to build:\n\nThen, a simplification pass removes nodes that are useless or that carry information which can otherwise be found implicitly in the graph:\n\nTo generate this representation as an HTML file (via `pyvis`\n\n), you can use the `cdg`\n\ntool provided by the `pcode_graph`\n\npackage:\n\n```\ncdg html --dataflow-only test.o -o test_dataflow.html\n```\n\nBut one can also directly generate the graph in mermaid format:\n\n```\ncdg md --dataflow-only test.o > test_dataflow.md\n```\n\nWe can notice on this graph that:\n\n- It is much simpler than the P-Code!\n- The permutation of instructions has no impact on the graph.\n\nThe Phi-node, however, leaves us puzzled: nothing tells us what makes EAX set to zero or to the sum of RDI and RSI. That's expected, since the data flow is not sufficient to express the semantics of the code. We therefore add control flow edges to the graph:\n\nIn the resulting graph, in addition to the condition for writing EAX, we also see the return address loading from the stack. The `EXTERNAL`\n\nnode expresses the jump to an address outside of the code present in the graph. If the code does not end with a branch, we will have an `END`\n\nnode. In total, we have ten different node kinds:\n\n```\nclass NodeKinds(Enum):\n\n    InputRegister = 0\n    OutputRegister = 1\n    Constant = 2\n\n    Operation = 3\n    Phi = 4\n\n    ReadMemory = 5\n    WrittenMemory = 6\n\n    Begin = 7\n    External = 8\n    End = 9\n```\n\nThere can be several memory nodes in a graph, but as no pointer aliasing analysis is performed we do not distinguish between memory addresses.\n\nYou may have noticed that adding the control flow edges breaks the invariance of the graph with respect to instruction permutations. We are thinking about improving the graph format to fix this. We are saving that for a future release of `pcode_graph`\n\n.\n\n## The CISCO TALOS dataset\n\nThis [USENIX2022 paper](https://www.usenix.org/system/files/sec22-marcelli.pdf) introduces [a dataset](https://github.com/Cisco-Talos/binary_function_similarity) including a large corpus of binaries compiled across:\n\n- 6 architectures (x86, ARM, MIPS)x(32 and 64 bits);\n- 8 compiler variants (4 for gcc and 4 for clang);\n- 5 optimization levels (O0, O1, O2, O3, Os).\n\nThe functions were extracted using IDA. In total, we have:\n\n- 256,625 functions for training;\n- 12,736 functions for validation;\n- 522,003 functions for testing.\n\nThey also publish a comparison of 10 state-of-the-art approaches on this dataset.\n\n## Extracting the graphs\n\nTo speed up the training of our models, we start with a preprocessing step to extract the CDG graphs corresponding to the 791,364 functions of the dataset.\n\nTo extract the graphs, we use the `LIEF`\n\nlibrary to extract the code at the offsets given by the dataset, then the `make_graph_from_binary`\n\nfunction to lift the code into P-Code and create the graph:\n\n``` python\nfrom pcode_graph.lief_importer import lookup_chunk\nfrom pcode_graph.maker import make_graph_from_binary\nfrom pcode_graph.translator import Translator\n\nfor arch, binaries in dataset_index.items():\n    translator = Translator(arch)\n\n    for binary_path, functions in binaries.items():\n\n        binary = parse_binary(binary_path)\n\n        for name, start, end in functions:\n            code = lookup_chunk(binary, start, end)\n\n            cdg = make_graph_from_binary(translator, code, start)\n\n            output_path = compute_graph_path(dataset_dir, binary_path, name)\n            output_path.write_bytes(pickle.dumps(cdg))\n```\n\nLifting binaries into P-Code with `pypcode`\n\nis trivial in a simple case, but properly handling the presence of data in the middle of the code, or errors in function extraction heuristics, is not so simple. You can have a look at the code of the `Translator`\n\nclass if you are interested.\n\nThe `make_graph_from_binary`\n\nmethod accepts options to guide the graph construction. In particular you can control the names of the registers whose writes should be considered as outputs of the code. If what you are extracting is a gadget, all registers can be useful, including processor flags. For our experiment, we used the default behavior which considers all the general purpose registers, but since we are interested in the semantics of a function, it would have been smarter to keep only the registers holding the return value, in order to simplify the dataflow graph while reducing the risk of bias.\n\nThis step takes time but parallelizes nicely via a small helper:\n\n```\ndef run_in_parallel[P, R](\n    function: Callable[[P], R],\n    parameters: list[P],\n    num_jobs: int | None = None,\n    initializer: Callable = lambda: None,\n) -> Iterator[R]:\n    with multiprocessing.Pool(processes=num_jobs, initializer=initializer) as pool:\n        for result in pool.imap_unordered(function, parameters):\n            yield result\n```\n\nTo avoid spending too much time on some very large functions of the dataset, we set a 5-second timeout for the analysis, the graph construction and its simplification. This removes 1.4% of the functions from the dataset.\n\nNote that trying to lift a binary into P-Code while specifying the wrong architecture can give surprising results, so we had to detect and skip the [316 labeling errors of the training set](https://github.com/Cisco-Talos/binary_function_similarity/issues/39).\n\n## Message-passing and diameter\n\nA graph neural network (GNN) learns a representation of each node through the so-called *message passing* technique, that is, by iteratively aggregating the information of its neighbors. Starting from an initial state, each node aggregates the states of its direct neighbors to update its own state.\n\nThis aggregation is done via a *graph convolution* operator. We use the `torch_geometric`\n\nlibrary, which provides an [implementation of a good part of the research on the topic](https://pytorch-geometric.readthedocs.io/en/latest/cheatsheet/gnn_cheatsheet.html):\n\nBy stacking *K* convolution layers, a node ends up integrating the information coming from its *K*-hop neighborhood, while taking the topological structure of the graph into account.\n\nThen, we can combine the resulting stats of all nodes to generate a global graph representation. This last stage is called *readout*.\n\nTo extract non-local properties of large graphs, one therefore has to stack many layers. However, unlike classical neural networks, GNNs suffer from an [oversmoothing](https://arxiv.org/pdf/2405.01663) problem when the number of hops increases (> 10 layers). This is one reason why we strive to produce graphs with the smallest possible diameter, independently of their size.\n\nTo get an idea, here is the distribution of graph diameters on the validation set when considering only the control-flow edges, only the data-flow edges, or both:\n\nWe notice that adding the control-flow to the graph slightly increases the diameter, but dataflow edges help limit the damage.\n\n## GNN architecture\n\nGNN architecture remains a vast research topic and searching for the best hyper-parameters takes time. We simply used [GINE](https://arxiv.org/abs/1905.12265) to have a solid baseline:\n\n``` python\nfrom torch import Tensor, relu\nimport torch\nfrom torch_geometric.data import Data\nfrom torch.nn import (\n    Dropout,\n    ReLU,\n    Linear,\n    Module,\n    ModuleList,\n    Sequential,\n)\nfrom torch_geometric.nn import GINEConv, global_add_pool, GraphNorm\nfrom dataclasses import dataclass\n\n@dataclass\nclass GNNConfig:\n    readout_head_outputs: int = 256\n    head_hidden: int = 256\n    conv_hidden: int = 64\n    conv_layers: int = 4\n    feature_dropout: float = 0.5\n\nclass GINE(Module):\n    def __init__(self, config: GNNConfig, node_features: int, edge_features: int):\n        super().__init__()\n        self.convs = ModuleList()\n        self.norms = ModuleList()\n        self.dropout = Dropout(config.feature_dropout)\n\n        for i in range(config.conv_layers):\n            dim_in = node_features if i == 0 else config.conv_hidden\n            self.convs.append(\n                GINEConv(\n                    Sequential(\n                        Linear(dim_in, config.conv_hidden),\n                        ReLU(),\n                        Linear(config.conv_hidden, config.conv_hidden),\n                    ),\n                    train_eps=True,\n                    edge_dim=edge_features,\n                )\n            )\n            self.norms.append(GraphNorm(config.conv_hidden))\n\n        self.head = Sequential(\n            Linear(config.conv_hidden * config.conv_layers, config.head_hidden),\n            ReLU(),\n            Dropout(config.feature_dropout),\n            Linear(config.head_hidden, config.readout_head_outputs),\n        )\n\n    def forward(self, data: Data) -> Tensor:\n        hs = []\n        x = data.x\n        for conv, norm in zip(self.convs, self.norms):\n            x = conv(x, data.edge_index, data.edge_attr)\n            x = relu(norm(x, data.batch))\n            x = self.dropout(x)\n            hs.append(global_add_pool(x, data.batch))\n\n        return self.head(torch.cat(hs, dim=-1))\n```\n\n## Loss function\n\nNow we have a model that computes an embedding of a graph, that is a fixed-size array of numbers associated with each graph of the batch passed to its `forward`\n\nmethod.\n\nBut it is not exactly what we need: to find similar functions we have to output a similarity score for a pair of functions.\n\nWe can do this by measuring the distance between embeddings, provided that the model was trained to push apart the embeddings of different functions while pulling together those of functions coming from the same program. Provided that our embeddings are L2-normalized, a dot is enough:\n\n```\nsimilarity = (emb1 * emb2).sum().item()\n```\n\nThe classical method to produce this kind of embeddings consists in using a Siamese model and a Triplet Margin Loss. We went for a technique presented at NeurIPS2020: the [Supervised Contrastive loss](https://proceedings.neurips.cc/paper_files/paper/2020/file/d89a66c7c80a29b1bdbab0f2a1a94af8-Paper.pdf) (SupCon loss).\n\nHere is the loss and the final model, which wraps the former one:\n\n``` python\nfrom torch import Tensor, matmul, eq, eye, exp, log, clamp\nfrom torch.nn import Module\nfrom torch_geometric.data import Data\nfrom torch.optim import AdamW\nfrom torch.nn.functional import normalize\n\ndef supcon_loss(features: Tensor, labels: Tensor, temperature: float):\n    \"\"\"\n    Supervised Contrastive Loss (Khosla et al., 2020).\n    \"\"\"\n\n    # Similarity matrix\n    # Should be already normalized\n    # features = normalize(features, dim=1)\n    logits = matmul(features, features.T) / temperature\n\n    # Hack for numeric stability\n    logits_max, _ = logits.max(dim=1, keepdim=True)\n    logits = logits - logits_max.detach()\n\n    # Compute mask of positive pairs (ie with same label)\n    labels = labels.view(-1, 1)\n    batch_size = features.shape[0]\n    positive_mask = eq(labels, labels.T).float()\n    self_mask = eye(batch_size, device=features.device)\n    positive_mask = positive_mask - self_mask\n\n    # Compute logprobs\n    logits_mask = 1.0 - self_mask\n    exp_logits = exp(logits) * logits_mask\n    log_prob = logits - log(exp_logits.sum(dim=1, keepdim=True) + 1e-12)\n\n    # Mean log_prob on positives for each anchor\n    num_positives = positive_mask.sum(dim=1)\n    mean_log_prob_pos = (positive_mask * log_prob).sum(dim=1) / clamp(\n        num_positives, min=1.0\n    )\n    loss = -mean_log_prob_pos.mean()\n    return loss\n\nclass SimilarityModel(Module):\n\n    def __init__(self, num_node_features: int, num_edge_features: int):\n        super().__init__()\n\n        self.gnn = GNN(num_node_features, num_edge_features)\n        self.optimizer = AdamW(self.gnn.parameters())\n\n    def forward(self, graph_batch: Data) -> Tensor:\n        \"\"\"Returns normalized embeddings of binaries with given graphs.\"\"\"\n\n        z = self.gnn(graph_batch)\n        return normalize(z, dim=1)\n\n    def step(self, batch):\n        emb = self(batch.graph)\n        loss = supcon_loss(emb, batch.func_id, 0.07)\n        self.optimizer.zero_grad()\n        loss.backward()\n        self.optimizer.step()\n```\n\n## Translating the graph into tensors\n\nThe model consumes graphs in the `torch.data.Dataset`\n\nformat. We still have to convert our `CDG`\n\ngraphs to this representation. For that we use the `graph_to_data`\n\nfunction provided by `pcode_graph`\n\n.\n\nThis method has many parameters to adapt the node features to the task at hand. The most important one controls the way registers are encoded. Since our task only concerns the semantics of the functions, there is no point in including the registers other than those used by the calling convention: whatever the registers assigned to variables, what matters is what is the semantics of the code.\n\nThe `map_calling_convention_registers`\n\nfunction of `pcode_graph`\n\nmakes it possible to encode the registers of the various calling conventions in a way that is consistent across architectures, so as to allow the model to generalize. It outputs an array of bits where each position corresponds to a use of the register: integer argument of rank n, 32-bit return value, etc. Note that a same register can be used both to pass a parameter and to return a value.\n\nIf your task only concerns a single architecture, you can simply use the `map_registers`\n\nfunction to convert a set of registers into hot-encoding. The idea is to send each possible value into a different input neuron, because it is much easier to teach a neural network a relation between its inputs than between different values of the same input.\n\nHere is the code, sparing you the loading part of the dataset:\n\n``` python\nimport pickle\nfrom pathlib import Path\nfrom typing import NamedTuple\nfrom torch import Tensor\nfrom torch.utils.data import Dataset\nfrom torch_geometric.data import Data\nfrom pcode_graph.gnn_exporter import graph_to_data, map_calling_convention_registers\n\nclass Function(NamedTuple):\n    bin_path: str\n    func_name: str\n    func_id: int\n    graph: Data\n\nclass FunctionDataset(Dataset):\n\n    def __init__(self, csv_path: Path):\n\n        super().__init__()\n\n        # Load dataset from CSV\n        self.graphs = []\n        architectures = set()\n        ...\n\n        # Create register mappings of same size to have the same amount of node features for each arch\n        self.register_mappers: dict[str, dict[str, Tensor]] = {}\n        for arch in architectures:\n            self.register_mappers[arch] = map_calling_convention_registers(arch)\n\n    def __len__(self) -> int:\n        return len(self.graphs)\n\n    def __getitem__(self, index) -> Function:\n        bin_path, arch, func_name, func_id, graph_path = self.graphs[index]\n        graph = pickle.loads(graph_path.read_bytes())\n        data = graph_to_data(graph, registers_emb=self.register_mappers[arch])\n        return Function(bin_path, func_name, func_id, data)\n\n    def __iter__(self):\n        for i in range(len(self)):\n            yield self[i]\n```\n\n## Building the batches\n\nThe SupCon loss needs several examples per class in order to be able to compare positive and negative pairs within a same batch. We write a *sampler* to build balanced batches containing several samples of several different functions:\n\n``` python\nfrom collections.abc import Iterator\nfrom itertools import islice\nfrom random import Random\nfrom torch.utils.data import Sampler\n\nSEED = 42\nSAMPLES_PER_FUNCTION = 4\nFUNCTIONS_PER_BATCH = 16\n\nclass BatchSampler(Sampler):\n    \"\"\"Dataset PK sampling to produce batch suitable for use with SupCon loss.\"\"\"\n\n    def __init__(self, dataset) -> None:\n        super().__init__()\n        self.dataset = dataset\n        self.rng = Random(SEED)\n        self.num_batches = sum(1 for _ in self)\n\n    def __len__(self) -> int:\n        return self.num_batches\n\n    def __iter__(self) -> Iterator[list[int]]:\n\n        # Collect shuffles samples\n        function_names = list(self.dataset.by_func_name.keys())\n        self.rng.shuffle(function_names)\n        samples: dict[str, list[int]] = {}\n\n        for func_name in function_names:\n            indices = list(self.dataset.by_func_name[func_name])\n            if len(indices) >= SAMPLES_PER_FUNCTION:\n                self.rng.shuffle(indices)\n                samples[func_name] = indices\n\n        # Build batches\n        while len(samples) >= FUNCTIONS_PER_BATCH:\n            batch = []\n            to_delete = []\n            for func_name, indexes in islice(samples.items(), FUNCTIONS_PER_BATCH):\n                batch += indexes[-SAMPLES_PER_FUNCTION:]\n                if len(indexes) >= 2 * SAMPLES_PER_FUNCTION:\n                    del indexes[-SAMPLES_PER_FUNCTION:]\n                else:\n                    to_delete.append(func_name)\n            for k in to_delete:\n                del samples[k]\n\n            yield batch\n```\n\n## Training loop\n\nIn general, the crux for a fast training is to be GPU-bound, meaning that the performance of your tool is tightly correlated to the performance of your GPU. This can be achieved provided that computations and transfers are overlapped. There are plenty of libraries doing that, but a small torch function with a ping-pong buffer does the trick:\n\n``` python\nfrom torch.utils.data import DataLoader\nfrom typing import Callable, NamedTuple\nfrom torch import Tensor\nfrom torch_geometric.data import Data\n\nclass BatchedFunctions(NamedTuple):\n    bin_path: list[str]\n    func_name: list[str]\n    func_id: Tensor\n    graph: Data\n\ndef apply_to_batches(\n    dataloader: DataLoader,\n    func: Callable[[BatchedFunctions], None],\n    device: str,\n):\n    \"\"\"Parallelizes computation and transfers with a ping-pong buffer,\n    at the cost of a higher GPU memory consumption.\n    \"\"\"\n\n    previous_batch: BatchedFunctions | None = None\n    batch: BatchedFunctions\n\n    for batch in dataloader:\n        next_batch = batch._replace(\n            graph=batch.graph.to(device, non_blocking=True),\n            func_id=batch.func_id.to(device, non_blocking=True),\n        )\n        if previous_batch is not None:\n            func(previous_batch)\n        previous_batch = next_batch\n\n    assert previous_batch, \"No batch returned by dataloader\"\n    func(previous_batch)\n```\n\nAnd to put all the pieces together, the training loop:\n\n``` python\nfrom torch_geometric.loader import DataLoader\n\ndef train(device, csv_path):\n\n    train_dataset = FunctionDataset(csv_path)\n    train_data = DataLoader(\n        dataset=train_dataset,\n        pin_memory=device == \"cuda\",\n        batch_sampler=BatchSampler(train_dataset),\n        num_workers=16,\n        persistent_workers=True,\n    )\n\n    model = SimilarityModel(train_dataset.node_features, train_dataset.edge_features)\n\n    model.to(device)\n    model.train()\n\n    for e in range(config.num_epochs):\n        apply_to_batches(train_data, model.step, device)\n```\n\nNote that the conversion from `Function`\n\nto `BatchedFunctions`\n\nis handled internally by the DataLoader.\n\nAlso note that to have a fully useful training pipeline you still have to incorporate:\n\n- Evaluation on the validation set (after each epoch for instance);\n- Logging (on tensorboard);\n- Best model saving.\n\n## Evaluation method and results\n\nThe authors provide a list of negative or positive pairs for the validation and test subsets. The algorithms to benchmark have to return a *similarity score* for each of these pairs.\n\nBy varying the threshold beyond which the tested functions are predicted as coming from the same program, one can plot a [ROC curve](https://en.wikipedia.org/wiki/Receiver_operating_characteristic). The metric used is the area under this curve.\n\nWe obtain the curves below. To compare algorithms, the benchmark provides several tasks, The hardest one, *XM*, consists in mixing all the possible variations:\n\nThe drawing of the distributions of similarity scores obtained on each kind of pairs shows that positive and negative pairs are clearly separated:\n\nAnd finally, we can compare our AUC score with the best one from the benchmark:\n\n| Model | XC (same arch and bitness) | XC+XB (same arch) | XA (same compiler) | XM (different arch, bitness, compiler & optim level) |\n|---|---|---|---|---|\n| GMN | 0.86 |\n0.87 |\n0.86 |\n0.87 |\n| GINE 4 layers + pcode_graph features | 0.86 |\n0.86 | 0.86 |\n0.87 |\n\nDoing the comparison is not completely honest because 1.5% of the bigger functions where just ignored in our case, but it seems that we achieve similar results that the best algorithm compared in the benchmark: GMN, for Graph Matching Networks.\n\nThis is not outstanding, but one could hardly expect more without working on the GNN architecture (and optimizing hyper-parameters at least a minimum).\n\n## Conclusion\n\nThis article shows how to use the `pcode_graph`\n\nlibrary to generate binary code embeddings based on semantics. Even using an old GNN baseline, we got promising results.\n\nWe could go further by exploring more recent GNN architectures (like GATv2, DirGNN...) and doing a decent hyper-parameter search (number and size of layers, dropout probability, kind of normalization, readout operators, etc.).\n\nDo not hesitate to fork the [GitHub repository](https://github.com/quarkslab/pcode_graph) to adapt the library to your own use-case and feel free to send us feedbacks or submit issues.", "url": "https://wpnews.pro/news/a-library-to-make-gnn-work-on-binary-code", "canonical_source": "https://blog.quarkslab.com/from-p-code-to-gnn-extract-binary-code-semantics.html", "published_at": "2026-08-20 09:59:02+00:00", "updated_at": "2026-08-20 10:15:16.403232+00:00", "lang": "en", "topics": ["machine-learning", "developer-tools"], "entities": ["Quarkslab", "pcode_graph", "Cisco-Talos", "ESANN 2026", "Sighthouse"], "alternates": {"html": "https://wpnews.pro/news/a-library-to-make-gnn-work-on-binary-code", "markdown": "https://wpnews.pro/news/a-library-to-make-gnn-work-on-binary-code.md", "text": "https://wpnews.pro/news/a-library-to-make-gnn-work-on-binary-code.txt", "jsonld": "https://wpnews.pro/news/a-library-to-make-gnn-work-on-binary-code.jsonld"}}