HOME       Features       Docs       ASPLOS 2023 Tutorial       GNN-Dataflow      Code       Resources       SST Integration       Flexagon

OMEGA GNN Cost Modeling Framework

OMEGA (Observing Mapping Efficiency over GNN Accelerator) framework is the cost model for inter-phase Graph Neural Network (GNN) dataflows. OMEGA can be used to model other SpMM and GEMM multiphase dataflows as well.

GNN Dataflow Exploration using OMEGA

GNNs are becoming increasingly popular because of their ability to accurately learn representations from graph structured data. GNN inference runtime is dominated by two phases: (1) Aggregation which is an SpMM computation with irregular, workload dependent data accesses, and (2) Combination computations that can be cast as GEMMs, similar to dense DNNs as shown in the figure above. Prior works on DNN dataflow studies have described the data orchestration and data movement in DNN accelerators. However, these works only model dense computations and model one GEMM or convolution operation at a time. GNNs offer an additional knob of pipelining between the two phases which also leads to interdependence of the two dataflows.

Taxonomy for GNN Dataflows

We aim to provide analysis of the design-space of GNN dataflows over flexible accelerator (for example - MAERI) which captures both individual phase dataflows (Intra-phase dataflows) and dataflows between the two phases (Inter-phase dataflows). To enable this, we propose a taxonomy that expresses: (1) Aggregation intra-phase dataflow (2) Combination intra-phase} dataflow (3) Inter-phase strategy, and (4) phase ordering targetting a flexible accelerator like MAERI which can support execution of all possible dataflows.

OMEGA Framework

We also demonstrate the OMEGA (Observing Mapping Efficiency over GNN Accelerator) framework that we build on top of STONNE which enables us to model the cost of the pipelined GNN dataflows. It instantiates SpMM and GEMM on STONNE’s flexible accelerator model MAERI and feeds the statistics to an inter-phase cost model that returns the metrics of a pipelined inter-phase dataflow as shown in Figure below.

Resources

Publications

For more details, please refer to our pre-print. Update: The paper has been accepted for publication in IPDPS 2022

Bibtex

If you use OMEGA and/or our GNN dataflow taxonomy in your reseach, please cite-

@inproceedings{garg2021understanding,
  title={Understanding the Design-Space of Sparse/Dense Multiphase GNN dataflows on Spatial Accelerators},
  author={Garg, Raveesh and Qin, Eric and Mu{\~n}oz-Mart{\'\i}nez, Francisco and Guirado, Robert and Jain, Akshay and Abadal, Sergi and Abell{\'a}n, Jos{\'e} L and Acacio, Manuel E and Alarc{\'o}n, Eduard and Rajamanickam, Sivasankaran and Krishna, Tushar},
  booktitle={2022 IEEE International Parallel and Distributed Processing Symposium (IPDPS)},
  year={2022}
  }

Codebase

The OMEGA codebase is available at https://github.com/stonne-simulator/omega and is also available in the docker image-

docker run -it franciscomunoz/stonne_omega_img /bin/bash