From fa752ae1e491a1f8fde8967bf04473c6a6c1ca18 Mon Sep 17 00:00:00 2001
From: Ian Wood <75152913+IanWood1@users.noreply.github.com>
Date: Tue, 29 Oct 2024 00:42:13 -0700
Subject: [PATCH] [DispatchCreation] Run preprocessing before elementwise
 fusion (#18920)

I think it makes sense to run `FusionPreprocessingPass` before
`ElementwiseOpFusionPass` because it helps put the IR in a better state
for fusion (e.g. interchanging `linalg.generic` indexing maps). But
also, reshapes have been propagated to the edges of the program, which
allows the `GatherFusionPattern` to be more effective.


Fixes compilation error from
https://github.com/iree-org/iree/issues/17226#issuecomment-2441200369.

---------

Signed-off-by: Ian Wood <ianwood2024@u.northwestern.edu>
---
 .../src/iree/compiler/DispatchCreation/Passes.cpp  | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/compiler/src/iree/compiler/DispatchCreation/Passes.cpp b/compiler/src/iree/compiler/DispatchCreation/Passes.cpp
index afee21cbbcd8..9cf5732962fd 100644
--- a/compiler/src/iree/compiler/DispatchCreation/Passes.cpp
+++ b/compiler/src/iree/compiler/DispatchCreation/Passes.cpp
@@ -127,9 +127,12 @@ static void addCleanupPatterns(OpPassManager &passManager) {
 //===----------------------------------------------------------------------===//
 
 void addDispatchRegionCreationPreprocessingPasses(OpPassManager &passManager) {
-  // 1. Do some simple elementwise op fusion. This could be skipped,
-  //    but could reduce the surface area of ops to handle later.
   FunctionLikeNest(passManager)
+      .addPass(IREE::Flow::createCanonicalizerPass)
+      .addPass(mlir::createCSEPass)
+      .addPass(DispatchCreation::createFusionPreprocessingPass)
+      // 1. Do some simple elementwise op fusion. This could be skipped,
+      //    but could reduce the surface area of ops to handle later.
       .addPass([]() {
         return DispatchCreation::createElementwiseOpFusionPass(
             ElementwiseOpFusionPassOptions{
@@ -148,6 +151,7 @@ void addDispatchRegionCreationPreprocessingPasses(OpPassManager &passManager) {
 
       // 3. Perform elementwise operation fusion again (now with higher
       //    dimensionality).
+      .addPass(DispatchCreation::createFusionPreprocessingPass)
       .addPass([]() {
         return DispatchCreation::createElementwiseOpFusionPass(
             ElementwiseOpFusionPassOptions{
@@ -294,12 +298,6 @@ void buildDispatchCreationPassPipeline(
         IREE::Util::createFixedPointIteratorPass(std::move(ipoPipeline)));
   }
 
-  FunctionLikeNest(passManager)
-      // Preprocess the input to a form more amenable for fusion.
-      .addPass(DispatchCreation::createFusionPreprocessingPass)
-      .addPass(IREE::Flow::createCanonicalizerPass)
-      .addPass(mlir::createCSEPass);
-
   addDispatchRegionCreationPreprocessingPasses(passManager);
   addDispatchRegionCreationPasses(passManager);