[mlir][Flow] Reland iree-org#14656 but keep CollapseReductionDims p…

…ass (iree-org#15219) iree-org#14656 attempted to keep a single collapse dimensions pass. But there are some cases that slip through that cause a segfault in the GPU backends (iree-org#15088) . Supporting those cases turn out to be actually involved. While figuring out, relanding the earlier commit to keep the gains that were made there.
nod-ai · Oct 19, 2023 · 479f4ed · 479f4ed
1 parent add9417
commit 479f4ed
Show file tree

Hide file tree

Showing 5 changed files with 561 additions and 193 deletions.
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel
@@ -96,6 +96,7 @@ iree_compiler_cc_library(
         "//llvm-external-projects/iree-dialects:IREELinalgTransformDialect",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AffineDialect",
+        "@llvm-project//mlir:AffineUtils",
         "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:ArithUtils",

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
@@ -78,6 +78,7 @@ iree_cc_library(
     IREELinalgTransformDialect
     LLVMSupport
     MLIRAffineDialect
+    MLIRAffineUtils
     MLIRAnalysis
     MLIRArithDialect
     MLIRArithUtils