TASO: optimizing deep learning computation with automatic generation of graph substitutions

[joapolarbear]

SOSP2020

[PDF] Related to 2019 SysML -- OPTIMIZING DNN COMPUTATION WITH RELAXED GRAPH SUBSTITUTIONS

[joapolarbear]

Other graph optimizers

TensorFlow, PyTorch, TensorRT, and TVM use a greedy rule-based optimization strategy and directly perform all applicable substitutions

What is Graph Substitution

1. A source graph (maybe a subgraph in a computation a graph)
2. A target graph, equivalent new subgraph
3. A mapping relation between input/output tensors in the source and target graph

Contribution

Graph Substitution Generator

The first DNN computation graph optimizer that automatically generates graph substitutions.

• Use a hash table to store the (graph, output), where the hash value is calculated based on the output, namely fingerprint.
• Input: a set of operators
• Traverse all possible permutations of operators with different input tensors (including the output tensors of previous operators --> construct a dependency)
• Using a depth-first search algorithm
• Hash those graphs that share the same set of output tensors (ignoring the permutation of output) to one fingerprinter.

Graph substitution Verifier

Check the correctness of all generated graph substitutions.

• Input: user-provided operator properties and graph substitutions.
• Test using floating-point numbers and keep the substitutions whose difference is smaller than a threshold
• • [ ] Check whether those properties can be met
• Pruning redundant substitutions: 1) input tensor renaming 2) remove common subgraph --> reduce the number of substitutions.

Joint Optimizer

• performs a cost-based backtracking search, applying the substitutions to find an optimized graph (similar to MetaFlow)

  considering layout transformations together with graph substitutions