Method and system for program sampling using neural network
This application describes methods, systems, and apparatus, for neural network-based program sampling (NPS). An example device may obtain an assembly code of a program and an execution trace of the program, and divide the assembly code into a plurality of execution intervals. The device may construct a plurality of code graphs respectively corresponding to the plurality of execution intervals, and for each of the plurality of code graphs: generate a plurality of graph snapshots based on the code graph and the execution trace of the program; embed, by using a Graph Neural Network, the plurality of graph snapshots into a plurality of vectors; and aggregate the plurality of vectors into an execution embedding. The device may cluster the plurality of execution embeddings into a plurality of clusters and select representative execution intervals of the program based on the plurality of clusters for execution.
1 . A computer-implemented method for program sampling, comprising:
obtaining an assembly code of a program and an execution trace of the program;
dividing the assembly code into a plurality of execution intervals, each execution interval corresponding to a portion of the assembly code;
constructing a plurality of code graphs respectively corresponding to the plurality of execution intervals;
for each code graph of the plurality of code graphs:
generating a plurality of graph snapshots based on the code graph and the execution trace of the program that correspond to different time points during an execution of the program, wherein generating each graph snapshot of the plurality of graph snapshots comprises fusing a subgraph of the code graph with dynamic states of the program from the execution trace of the program, and the fusing comprises adding or removing one or more nodes or edges of the code graph based on the dynamic states of the program at a given time point during the execution of the program;
embedding, by using a Graph Neural Network (GNN), the plurality of graph snapshots into a plurality of vectors; and
aggregating the plurality of vectors corresponding to the plurality of graph snapshots into an execution embedding to represent the execution interval corresponding to the code graph;
clustering the plurality of execution embeddings representing the plurality of execution intervals into a plurality of clusters; and
selecting representative execution intervals of the program based on the plurality of clusters for execution.
2 . The computer-implemented method of claim 1 , wherein the obtaining the assembly code of the program comprises:
compiling the program written in a high-level programming language into the assembly code of the program.
3 . The computer-implemented method of claim 1 , wherein the constructing the plurality of code graphs respectively corresponding to the plurality of execution intervals comprises, for each execution interval of the plurality of execution intervals:
identifying instructions and variables in assembly code in the execution interval;
constructing a backbone of a code graph by:
creating instruction nodes and variable nodes based on the instructions and variables;
connecting the instruction nodes and the variable nodes to represent a control flow and a data flow of the assembly code;
adding function nodes to the backbone of the code graph to represent functions that are implied in the assembly code; and
connecting the instruction nodes, variable nodes, and function nodes to form the code graph.
4 . The computer-implemented method of claim 3 , wherein the function nodes correspond to functions including memory reference functions.
5 . The computer-implemented method of claim 1 , wherein the generating the plurality of graph snapshots based on the code graph and the execution trace of the program comprises:
for a given graph node in the code graph, determining a memory depth based on a number of memory references visited along all code paths in the code graph that start from the given graph node;
performing a breadth-first search in the code graph starting from the given graph node constrained by the memory depth to obtain a subgraph of the code graph; and
adding intermediate nodes representing the dynamic states extracted from the execution trace to the subgraph to obtain the graph snapshot.
6 . The computer-implemented method of claim 1 , wherein the GNN is a graph attention network (GAT) comprising parameters jointly trained with a code path selection module and a memory address prediction module.
7 . The computer-implemented method of claim 6 , wherein the GAT is trained by:
feeding a training graph snapshot into the GAT for embedding, wherein the training graph snapshot corresponds to a ground truth execution trace;
obtaining, from the GAT, embeddings of function nodes in the training graph snapshot, wherein the function nodes correspond to memory reference functions in the training graph snapshot;
feeding the embeddings of the function nodes into the code path selection module to predict a code path in the training graph snapshots;
obtaining embeddings of memory addresses to be visited along the code path;
feeding the embeddings of the memory addresses into the memory address prediction module to predict next memory addresses to be visited along the code path;
obtaining prediction errors by comparing the predicted next memory addresses and the ground truth execution trace; and
adjusting parameters of the GAT, the code path selection module, and the memory address prediction module to minimize the prediction errors.
8 . The computer-implemented method of claim 6 , wherein the code path selection module comprises an attention layer and a weighted unsort segment sum layer for generating weights for code path candidates, wherein a code path candidate with a highest weight is selected as a predicted code path.
9 . The computer-implemented method of claim 6 , wherein after training, the GAT is deployed for embedding graph snapshots.
10 . The computer-implemented method of claim 1 , wherein clustering the plurality of execution embeddings of the plurality of execution intervals into the plurality of clusters comprises:
clustering the plurality of execution embeddings using K-means.
11 . The computer-implemented method of claim 1 , wherein the selecting representative execution intervals of the program based on the plurality of clusters comprises:
from one or more of the plurality of clusters, selecting one execution embedding from each of the plurality of clusters; and
selecting, from the program, one or more execution intervals that correspond to the one or more selected execution embeddings as the representative execution intervals.
12 . The computer-implemented method of claim 1 , further comprising:
executing the representative execution intervals of the program rather than executing an entirety of the program.
13 . A system comprising one or more processors and one or more non-transitory computer-readable memories coupled to the one or more processors, the one or more non-transitory computer-readable memories storing instructions that, when executed by the one or more processors, cause the system to perform operations comprising:
obtaining an assembly code of a program and an execution trace of the program;
dividing the assembly code into a plurality of execution intervals, each execution interval corresponding to a portion of the assembly code;
constructing a plurality of code graphs respectively corresponding to the plurality of execution intervals;
for each code graph of the plurality of code graphs:
generating a plurality of graph snapshots based on the code graph and the execution trace of the program that correspond to different time points during an execution of the program, wherein generating each graph snapshot of the plurality of graph snapshots comprises fusing a subgraph of the code graph with dynamic states of the program from the execution trace of the program, and the fusing comprises adding or removing one or more nodes or edges of the code graph based on the dynamic states of the program at a given time point during the execution of the program;
embedding, by using a Graph Neural Network (GNN), the plurality of graph snapshots into a plurality of vectors; and
aggregating the plurality of vectors corresponding to the plurality of graph snapshots into an execution embedding to represent the execution interval corresponding to the code graph;
clustering the plurality of execution embeddings representing the plurality of execution intervals into a plurality of clusters; and
selecting representative execution intervals of the program based on the plurality of clusters for execution.
14 . The system of claim 13 , wherein the generating the plurality of graph snapshots based on the code graph and the execution trace of the program comprises:
for a given graph node in the code graph, determining a memory depth based on a number of memory references visited along all code paths in the code graph that start from the given graph node;
performing a breadth-first search in the code graph starting from the given graph node constrained by the memory depth to obtain a subgraph of the code graph; and
adding intermediate nodes representing the dynamic states extracted from the execution trace to the subgraph to obtain the graph snapshot.
15 . The system of claim 13 , wherein the GNN is a graph attention network (GAT) comprising parameters jointly trained with a code path selection module and a memory address prediction module.
16 . The system of claim 15 , wherein the GAT is trained by:
feeding a training graph snapshot into the GAT for embedding, wherein the training graph snapshot corresponds to a ground truth execution trace;
obtaining, from the GAT, embeddings of function nodes in the training graph snapshot, wherein the function nodes correspond to memory reference functions in the training graph snapshot;
feeding the embeddings of the function nodes into the code path selection module to predict a code path in the training graph snapshots;
obtaining embeddings of memory addresses to be visited along the code path;
feeding the embeddings of the memory addresses into the memory address prediction module to predict next memory addresses to be visited along the code path;
obtaining prediction errors by comparing the predicted next memory addresses and the ground truth execution trace; and
adjusting parameters of the GAT, the code path selection module, and the memory address prediction module to minimize the prediction errors.
17 . The system of claim 15 , wherein the code path selection module comprises an attention layer and a weighted unsort segment sum layer for generating weights for code path candidates, wherein a code path candidate with a highest weight is selected as a predicted code path.
18 . The system of claim 13 , wherein the selecting representative execution intervals of the program based on the plurality of clusters comprises:
from one or more of the plurality of clusters, selecting one execution embedding from each of the plurality of clusters; and
selecting, from the program, one or more execution intervals that correspond to the one or more selected execution embeddings as the representative execution intervals.
19 . The system of claim 13 , wherein the operations further comprise:
executing the representative execution intervals of the program rather than executing an entirety of the program.
20 . A non-transitory computer-readable storage medium configured with instructions executable by one or more processors to cause the one or more processors to perform operations comprising:
obtaining an assembly code of a program and an execution trace of the program;
dividing the assembly code into a plurality of execution intervals, each execution interval corresponding to a portion of the assembly code;
constructing a plurality of code graphs respectively corresponding to the plurality of execution intervals;
for each code graph of the plurality of code graphs:
generating a plurality of graph snapshots based on the code graph and the execution trace of the program that correspond to different time points during an execution of the program, wherein generating each graph snapshot of the plurality of graph snapshots comprises fusing a subgraph of the code graph with dynamic states of the program from the execution trace of the program, and the fusing comprises adding or removing one or more nodes or edges of the code graph based on the dynamic states of the program at a given time point during the execution of the program;
embedding, by using a Graph Neural Network (GNN), the plurality of graph snapshots into a plurality of vectors; and
aggregating the plurality of vectors corresponding to the plurality of graph snapshots into an execution embedding to represent the execution interval corresponding to the code graph;
clustering the plurality of execution embeddings representing the plurality of execution intervals into a plurality of clusters; and
selecting representative execution intervals of the program based on the plurality of clusters for execution.