[SOSP'19] TASO: Optimizing Deep Learning Computation with Automated Generation of Graph Substitutions

[ganler]

https://cs.stanford.edu/~zhihao/papers/sosp19.pdf

---

3rdparty review on ZHIHU(Chinese): https://zhuanlan.zhihu.com/p/94613279

[ganler]

current ml systems' approach

Rule-based operator fusion.

TensorFlow XLA: 200+ rules and 53'000+ LOC.

Limitations

• Robustness: not apply all models/hardware.
• Scalability: New operators and graph structures req. more rules.
• Performance: Miss subtle opt. for specific scenarios.

(Credit: https://www.youtube.com/watch?v=YsJICemFBsQ&t=234s)

TASO

High-level Arch

Generate => Approximation-based verification/filtering => Mathematic verification

Tasks

• Generation of potential substitutions.
• Verification.

Graph Generation

• Enumerate all possible combinations of available operators.
• Compute output fingerprints with random input tensors.

Graph Verifier

743 subs. in 5 minutes, and verifies them on 43 operators in 10 minutes. Zhihao is also working on distributed fingerprint generation to push the limit of the optimizing operator numbers.

Search-Based Graph Optimizer

• Cost model.
• Search-Based Graph Optimizer

Results

Possible Questions / Comments

• I feel this work is very useful in production, and graph-level optimization is more on-the-fly and consistently working for most cases.
• There can be more challenges in this topic:
  • Learning optimization experiences. (Can we built an OptimizationNet(like ImageNet) to learn the optimization worldwide) If we can, how can we accelerate them? Also for security, we may need to send profiling results over the cloud.
  • For device placement.
  • Like Zhihao mentioned, increase the operator number in graph substitution generation, and go distributed.

Fingerprint:

• If the tensor program only has a very limited discrete output range, it may fail on correctness.
• Numerical limits changed... (But of course, without the change, it's hard to optimize tensor programs)
  • (a/c + b/c) has bigger numeric range than (a + b)/c. (The latter is easier to overflow)

Comparison with related works

TVM | Tensor Comprehensions

• Focusing on optimizing individual operators.
• Rely on manually defined search space.