PET: Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections

https://www.usenix.org/conference/osdi21/presentation/wang

• Current system considers only full equivalent transformations

  • Pro: preserve functionality

  • Con: miss optimization opportunities

• Partial Equivalent Transformation

  • Pro: better performance

    • Faster ML operators

    • More efficient tensor layouts

    • HW-specific optimizations

  • Con: potential accuracy loss

• Key: benefits while preserving equivalence?

  • Correction preserves equivalence

1. Super optimization

2. Multi-linearity of DNN computations

   1. DNN computation = MLTP + non-linear activations

   2. Reduce complexity from O(m*n) to O(1)

Corrector:

1. Re-compute the incorrect outputs using the original program

2. opportunistically fuse correction kernels with other operators

3. Eval: less than 1% overhead

Program optimizer:

1. Search-based program optimizer

   1. Beam search