| 624 | #include "tensorflow/compiler/mlir/lite/transforms/generated_prepare_tf.inc" |
| 625 | |
| 626 | void PrepareTFPass::runOnFunction() { |
| 627 | OwningRewritePatternList patterns; |
| 628 | auto func = getFunction(); |
| 629 | MLIRContext *ctx = &getContext(); |
| 630 | |
| 631 | // This pattern was intented to uses TFL QDQs to preserve the quantization |
| 632 | // parameters from the TF Quant ops, thus this pattern should run with the |
| 633 | // first `applyPatternsGreedily` method, which would otherwise removes the |
| 634 | // TF FakeQuant ops by the constant folding. |
| 635 | patterns.insert<PreparePerTensorFakeQuant, PreparePerChannelFakeQuant>(ctx); |
| 636 | |
| 637 | // This pattern will try to identify and optimize for dilated convolution. |
| 638 | // e.g. Patterns like "SpaceToBatchND -> Conv2D -> BatchToSpaceND" will be |
| 639 | // replaced with a single Conv op with dilation parameter. |
| 640 | patterns.insert<ConvertTFDilatedConvOp<TF::Conv2DOp>, |
| 641 | ConvertTFDilatedConvOp<TF::DepthwiseConv2dNativeOp>>(ctx); |
| 642 | TFL::populateWithGenerated(ctx, &patterns); |
| 643 | // TODO(karimnosseir): Split to separate pass probably after |
| 644 | // deciding on long term plan for this optimization. |
| 645 | // This will allow optimizing any TF_Mul->TF_Conv in the graph |
| 646 | // and any expanded from FusedBatchNorm. We need to do this |
| 647 | // before converting TF_Conv to TFL_Conv |
| 648 | applyPatternsGreedily(func, patterns); |
| 649 | |
| 650 | // Load the generated pattern again, so new quantization pass-through |
| 651 | // will be applied. |
| 652 | patterns.clear(); |
| 653 | TFL::populateWithGenerated(ctx, &patterns); |
| 654 | if (unfold_batch_matmul_) { |
| 655 | patterns.insert<TF::ConvertTFBatchMatMulOp<TF::BatchMatMulOp>, |
| 656 | TF::ConvertTFBatchMatMulOp<TF::BatchMatMulV2Op>>(ctx); |
| 657 | } |
| 658 | patterns.insert<ConvertTFConv2D, ConvertTFDepthwiseConv2dNative, |
| 659 | ConvertTFStridedSlice>(ctx); |
| 660 | applyPatternsGreedily(func, patterns); |
| 661 | } |
| 662 | |
| 663 | } // namespace |
| 664 | |