  Tensor graph superoptimisation systems perform a sequence of subgraph
substitution to neural networks, to find the optimal computation graph
structure. Such a graph transformation process naturally falls into the
framework of sequential decision-making, and existing systems typically employ
a greedy search approach, which cannot explore the whole search space as it
cannot tolerate a temporary loss of performance. In this paper, we address the
tensor graph superoptimisation problem by exploring an alternative search
approach, reinforcement learning (RL). Our proposed approach, X-RLflow, can
learn to perform neural network dataflow graph rewriting, which substitutes a
subgraph one at a time. X-RLflow is based on a model-free RL agent that uses a
graph neural network (GNN) to encode the target computation graph and outputs a
transformed computation graph iteratively. We show that our approach can
outperform state-of-the-art superoptimisation systems over a range of deep
learning models and achieve by up to 40% on those that are based on
transformer-style architectures.
