From 7feb0b28ad5045e9d5c6fcf528bbc048f8969381 Mon Sep 17 00:00:00 2001
From: zjgarvey <47986913+zjgarvey@users.noreply.github.com>
Date: Tue, 16 Jul 2024 10:33:18 -0700
Subject: [PATCH] Add torch-fuse-quantized-ops pass to the torch-to-iree
 pipeline (#17908)

The torch to iree pipeline currently does not use
`--torch-fuse-quantized-ops`, which is the cause of significant
discrepancies between model testing with iree-compile from torch IR and
model testing which first lowers to linalg with torch-mlir before
compiling. Together with `--torch-fuse-quantized-ops`, a newer pass
`--torch-scalarize-shapes` is added to the `torch-to-iree` pipeline to
keep in line with the
`--torch-backend-to-linalg-on-tensors-backend-pipeline`.

---------

Signed-off-by: zjgarvey <zjgarvey@gmail.com>
Signed-off-by: Lubo Litchev <lubol@google.com>
---
 compiler/plugins/input/Torch/InputConversion/Passes.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/compiler/plugins/input/Torch/InputConversion/Passes.cpp b/compiler/plugins/input/Torch/InputConversion/Passes.cpp
index 2dc4a9291d0c6..293921892b26d 100644
--- a/compiler/plugins/input/Torch/InputConversion/Passes.cpp
+++ b/compiler/plugins/input/Torch/InputConversion/Passes.cpp
@@ -49,6 +49,8 @@ void createTorchToIREEPipeline(
       mlir::torch::TorchConversion::createConvertCustomQuantOpPass());
   pm.addNestedPass<func::FuncOp>(
       torch::Torch::createDecomposeComplexOpsPass(emptyArrayRef));
+  pm.addNestedPass<func::FuncOp>(torch::Torch::createFuseQuantizedOpsPass());
+  pm.addNestedPass<func::FuncOp>(torch::Torch::createScalarizeShapesPass());
   pm.addNestedPass<func::FuncOp>(torch::createConvertTorchToTMTensorPass());
   pm.addNestedPass<func::FuncOp>(
       TorchInput::createConvertTMTensorToLinalgExtPass());