Graph neural networks with configuration cross-attention for tensor compilers
<|reference_start|>Graph neural networks with configuration cross-attention for tensor compilers: With the recent popularity of neural networks comes the need for efficient serving of inference workloads. A neural network inference workload can be represented as a computational graph with nodes as operators transforming multidimensional tensors. The tensors can be transposed and/or tiled in a combinatorially large number of ways, some configurations leading to accelerated inference. We propose TGraph, a neural graph architecture that allows screening for fast configurations of the target computational graph, thus representing an artificial intelligence (AI) tensor compiler in contrast to the traditional heuristics-based compilers. The proposed solution improves mean Kendall's $\tau$ across layout collections of TpuGraphs from 29.8% of the reliable baseline to 67.4% of TGraph. We estimate the potential CO$_2$ emission reduction associated with our work to be equivalent to over 50% of the total household emissions in the areas hosting AI-oriented data centers.<|reference_end|>
@article{khizbullin2024graph,
    title={Graph neural networks with configuration cross-attention for tensor
  compilers},
    author={Dmitrii Khizbullin, Eduardo Rocha de Andrade, Thanh Hau Nguyen,
  Matheus Pedroza Ferreira and David R. Pugh},
    journal={arXiv preprint arXiv:2405.16623},
    year={2024},
    archivePrefix={arXiv},
    eprint={2405.16623},
    primaryClass={cs.LG cs.AR cs.PF}
}