Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Layout transform: Fix-up QDQ units and add constant folding (microsof…
…t#20685) ### Description #### Problem 1: Broken Transpose QDQ unit Layout transform's specialized cost function aggressively pushes down transposes with channel-first or channel-last perms. This can lead to a situation where a channel-fist/last Transpose gets stuck after being pushed through an Unsqueeze node that makes the Transpose's perm no longer channel-first/last. At this point, the specialized cost function defers to the default const function, which does not see a need to continue pushing this transpose node. This breaks the QDQ node units for both the Unsqueeze and the Transpose: DQ -> Unsqueeze -> Transpose -> Q. <img width="266" alt="image" src="https://github.com/microsoft/onnxruntime/assets/19691973/82f8432d-ca27-451b-8c36-c8d87b806e30"> The transpose optimizer should insert a Q -> DQ pair between the Unsqueeze and Transpose nodes to fix both QDQ node units: DQ -> Unsqueeze -> Q[new] -> DQ[new] -> Transpose -> Q <img width="198" alt="image" src="https://github.com/microsoft/onnxruntime/assets/19691973/5a584bdf-e5db-4622-b3bb-83c060e09261"> #### Problem 2: Inserted Squeeze/Transpose nodes should be constant folded when possible. The transpose optimizer inserts Squeeze (and Transpose) ops between an initializer and a DQ to counteract the effect of Unsqueezing that initializer if it is consumed by multiple nodes. This results in a graph where the inserted nodes are not in valid node units: Original graph where two Mul nodes share a common initializer input: <img width="456" alt="image" src="https://github.com/microsoft/onnxruntime/assets/19691973/4b9155ae-e32f-41fc-9136-f953b73e92e7"> Resulting graph after transpose optimization without constant folding: <img width="452" alt="image" src="https://github.com/microsoft/onnxruntime/assets/19691973/3c1bfef1-d45f-4d6e-aa19-1c2929eae3f5"> Here, the circled Transpose and Squeeze nodes operate on a quantized integer type but are not in valid QDQ node units. The solution is to run constant folding, which results in: <img width="405" alt="image" src="https://github.com/microsoft/onnxruntime/assets/19691973/aebdb91f-f38f-4583-adec-33e46126365f"> ### Motivation and Context Improve the layout transformation to allow more models to run on EPs that prefer the channel-last layout. --------- Co-authored-by: Scott McKay <[email protected]>
- Loading branch information