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

  • 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)


  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

Last updated