Move more pass from Flow stage to GlobalOptimization stage. (#14707)

- Move four more passes to GlobalOptimization stage. - ConvertElementwiseToLinalgPass - GeneralizeLinalgNamedOpsPass - FuseDequantizationMatmulPass - FoldUnitExtentDimsPass - Move Flow transformation_pipeline.mlir test to GlobalOptimization/test. It is mainly for testing ConvertElementwiseToLinalg pass which is also tested upstream. We probably can remove it as a follow-up.
iree-org · Aug 28, 2023 · afa74de · afa74de
1 parent b890e22
commit afa74de
Show file tree

Hide file tree

Showing 12 changed files with 94 additions and 42 deletions.
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
@@ -132,12 +132,6 @@ void buildFlowTransformPassPipeline(OpPassManager &passManager,
 
   FunctionLikeNest(passManager)
       // Preprocess the input to a form more amenable for fusion
-      // - Convert all elementwise ops to Linalg
-      // - Remove unit-extent dimensions.
-      .addPass(mlir::createConvertElementwiseToLinalgPass)
-      .addPass(createGeneralizeLinalgNamedOpsPass)
-      .addPass(createFuseDequantizationMatmulPass)
-      .addPass(createFoldUnitExtentDimsPass)
       .addPass(createRaiseSpecialOps)
       .addPass(createInterchangeGenericOpsPass)
       .addPass(createCollapseDimsPass)

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel
@@ -57,7 +57,6 @@ iree_lit_test_suite(
             "tensor_pad_to_tensor_insert_slice.mlir",
             "top_level_scf_to_cfg.mlir",
             "transform_dispatch_region_formation.mlir",
-            "transformation_pipeline.mlir",
             "verify_input_ir.mlir",
         ],
         include = ["*.mlir"],

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
@@ -55,7 +55,6 @@ iree_lit_test_suite(
     "tensor_pad_to_tensor_insert_slice.mlir"
     "top_level_scf_to_cfg.mlir"
     "transform_dispatch_region_formation.mlir"
-    "transformation_pipeline.mlir"
     "verify_input_ir.mlir"
   TOOLS
     FileCheck

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/pipeline_tests.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/pipeline_tests.mlir
@@ -1,4 +1,5 @@
-// RUN: iree-opt --iree-flow-transformation-pipeline --split-input-file %s | FileCheck %s
+// TODO(hanchung): Split the transformation pipeline tests into two mlir files.
+// RUN: iree-opt --iree-global-optimization-transformation-pipeline --iree-flow-transformation-pipeline --split-input-file %s | FileCheck %s
 
 #map = affine_map<(d0, d1) -> (d0)>
 #map1 = affine_map<(d0, d1) -> (d1)>

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/transformation_pipeline.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/transformation_pipeline.mlir
diff --git a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
@@ -29,6 +29,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:LinalgTransforms",
+        "@llvm-project//mlir:MemRefTransforms",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Transforms",
     ],

diff --git a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
@@ -22,6 +22,7 @@ iree_cc_library(
     MLIRFuncDialect
     MLIRIR
     MLIRLinalgTransforms
+    MLIRMemRefTransforms
     MLIRPass
     MLIRTransforms
     iree::compiler::Dialect::Flow::Transforms

diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
@@ -10,6 +10,7 @@
 #include "iree/compiler/Utils/PassUtils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Transforms/Passes.h"
 
 namespace mlir {
@@ -51,12 +52,23 @@ void buildGlobalOptimizationPassPipeline(
   mainPassManager.addPass(IREE::Flow::createEraseUnusedLinalgOperands());
 
   // Expand tensor shapes into SSA values and optimize the whole program.
-  // The more we are able to equate shape dimensions at this level the better
-  // our fusions will be.
+  // The more we are able to equate shape dimensions at this level the
+  // better our fusions will be.
   FunctionLikeNest(mainPassManager)
       .addPass(IREE::Flow::createTopLevelSCFToCFGPass);
   mainPassManager.addPass(IREE::Flow::createExpandTensorShapesPass());
 
+  FunctionLikeNest(mainPassManager)
+      // Preprocess the input to a form more amenable for fusion
+      // - Convert all elementwise ops to Linalg
+      // - Remove unit-extent dimensions.
+      .addPass(mlir::createConvertElementwiseToLinalgPass)
+      .addPass(IREE::Flow::createGeneralizeLinalgNamedOpsPass)
+      .addPass(IREE::Flow::createFuseDequantizationMatmulPass)
+      .addPass(IREE::Flow::createFoldUnitExtentDimsPass)
+      .addPass(mlir::createCanonicalizerPass)
+      .addPass(mlir::createCSEPass);
+
   OpPassManager pipeline(ModuleOp::getOperationName());
   FunctionLikeNest(pipeline)
       // Simplify util.global accesses early on; this can help with dispatch

diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel
@@ -0,0 +1,28 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
+
+package(
+    features = ["layering_check"],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+iree_lit_test_suite(
+    name = "lit",
+    srcs = enforce_glob(
+        [
+            "transformation_pipeline.mlir",
+        ],
+        include = ["*.mlir"],
+    ),
+    cfg = "//compiler:lit.cfg.py",
+    tools = [
+        "//tools:iree-opt",
+        "@llvm-project//llvm:FileCheck",
+    ],
+)
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt
@@ -0,0 +1,23 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from           #
+# compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel               #
+#                                                                              #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary   #
+# CMake-only content.                                                          #
+#                                                                              #
+# To disable autogeneration for this file entirely, delete this header.        #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_lit_test_suite(
+  NAME
+    lit
+  SRCS
+    "transformation_pipeline.mlir"
+  TOOLS
+    FileCheck
+    iree-opt
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/transformation_pipeline.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/transformation_pipeline.mlir
@@ -0,0 +1,24 @@
+// RUN: iree-opt --split-input-file --iree-global-optimization-transformation-pipeline %s | FileCheck %s
+
+// CHECK-LABEL: @empty
+func.func @empty() {
+  // CHECK-NEXT: return
+  return
+}
+
+// -----
+
+func.func @elementwiseOps(%arg0 : tensor<4xf32>) -> tensor<4xf32> {
+  %0 = arith.addf %arg0, %arg0 : tensor<4xf32>
+  %1 = arith.subf %0, %arg0 : tensor<4xf32>
+  %2 = arith.mulf %1, %arg0 : tensor<4xf32>
+  return %2 : tensor<4xf32>
+}
+
+// CHECK-LABEL: func.func @elementwiseOps(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+//       CHECK:   %{{.+}} = linalg.generic
+//       CHECK:     %{{.+}} = arith.addf %{{.+}}, %{{.+}} : f32
+//       CHECK:   %{{.+}} = linalg.generic
+//       CHECK:     %{{.+}} = arith.subf %{{.+}}, %{{.+}} : f32
+//       CHECK:   %{{.+}} = linalg.generic
+//       CHECK:     %{{.+}} = arith.mulf %{{.+}}, %{{.+}} : f32
diff --git a/tools/test/compile_pipelines.mlir b/tools/test/compile_pipelines.mlir
@@ -1,6 +1,7 @@
 // RUN: iree-opt --iree-common-input-transformation-pipeline %s | \
 // RUN: iree-opt --iree-abi-transformation-pipeline - | \
 // RUN: iree-opt --iree-common-input-transformation-pipeline - | \
+// RUN: iree-opt --iree-global-optimization-transformation-pipeline - | \
 // RUN: iree-opt --iree-flow-transformation-pipeline - | \
 // RUN: iree-opt --iree-stream-transformation-pipeline - | \
 // RUN: iree-opt --iree-hal-transformation-pipeline --iree-hal-target-backends=vmvx - | \