Revert "[GPU][Codegen] Expand iteration space based on new expand_dims attribute (#22342) (#22982)

This reverts commit b879ee1.

To temporarily fix ONNX models test suite from issues suspected from
[[DispatchCreation] Unit dims not properly folded away during
GlobalOptimization](#22978) and
reported regression.

Signed-off-by: Eric Feng <[email protected]>

Author: efric

compiler/src/iree/compiler/Codegen/Common/BlockDynamicDimensions.cpp

@@ -36,6 +36,17 @@ using TensorDivisibilityInfo =
 
 namespace {
 
+struct RemoveOptimizationBarrier final
+    : public OpRewritePattern<IREE::Util::OptimizationBarrierOp> {
+  using Base::Base;
+
+  LogicalResult matchAndRewrite(IREE::Util::OptimizationBarrierOp barrierOp,
+                                PatternRewriter &rewriter) const override {
+    rewriter.replaceOp(barrierOp, barrierOp.getOperands());
+    return success();
+  }
+};
+
 /// This pass is used to materialize information about dynamic dimensions of
 /// `tensor` operands of an operation in the IR. If a dynamic dimension is
 /// known to be a multiple of a compile-time constant value, this pass
@@ -99,6 +110,10 @@ getTensorDivisibilityInfo(const TensorDynamicDimAnalysis &dynamicDimAnalysis,
 /// inverses of each other. The `util.optimization.barrier` avoid these from
 /// getting folded away during reshape propagation. Return the result of the
 /// `tensor.collapse_shape generated.
+struct ReshapeOps {
+  tensor::ExpandShapeOp expandShapeOp;
+  tensor::CollapseShapeOp collapseShapeOp;
+};
 static std::optional<ReshapeOps>
 blockDynamicDimensionsOfValue(RewriterBase &rewriter,
                               const TensorDivisibilityInfo &divisibilityInfo,
@@ -398,7 +413,7 @@ void BlockDynamicDimensionsPass::runOnOperation() {
   // Delete the optimization barrier and run some further cleanup.
   {
     RewritePatternSet removeBarrierOpsPatterns(context);
-    populateRemoveOptimizationBarrierPatterns(removeBarrierOpsPatterns);
+    removeBarrierOpsPatterns.insert<RemoveOptimizationBarrier>(context);
     tensor::ExpandShapeOp::getCanonicalizationPatterns(removeBarrierOpsPatterns,
                                                        context);
     tensor::CollapseShapeOp::getCanonicalizationPatterns(

compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel

@@ -74,7 +74,6 @@ iree_compiler_cc_library(
         "GPUDistributeScfFor.cpp",
         "GPUDistributeSharedMemoryCopy.cpp",
         "GPUDistributionPatterns.cpp",
-        "GPUExpandDimensions.cpp",
         "GPUFuseAndHoistParallelLoops.cpp",
         "GPUGeneralizeNamedOps.cpp",
         "GPUGreedilyDistributeToThreads.cpp",
@@ -126,7 +125,6 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Dialect/LinalgExt/Transforms",
         "//compiler/src/iree/compiler/Dialect/LinalgExt/Utils",
         "//compiler/src/iree/compiler/Dialect/TensorExt/IR",
-        "//compiler/src/iree/compiler/Dialect/Util/IR",
         "//compiler/src/iree/compiler/Utils",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AMDGPUDialect",

compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt

@@ -67,7 +67,6 @@ iree_cc_library(
     "GPUDistributeScfFor.cpp"
     "GPUDistributeSharedMemoryCopy.cpp"
     "GPUDistributionPatterns.cpp"
-    "GPUExpandDimensions.cpp"
     "GPUFuseAndHoistParallelLoops.cpp"
     "GPUGeneralizeNamedOps.cpp"
     "GPUGreedilyDistributeToThreads.cpp"
@@ -160,7 +159,6 @@ iree_cc_library(
     iree::compiler::Dialect::LinalgExt::Transforms
     iree::compiler::Dialect::LinalgExt::Utils
     iree::compiler::Dialect::TensorExt::IR
-    iree::compiler::Dialect::Util::IR
     iree::compiler::Utils
   PUBLIC
 )

compiler/src/iree/compiler/Codegen/Common/GPU/GPUExpandDimensions.cpp

@@ -383,14 +383,6 @@ def GPUApplyPaddingLevelPass :
   ];
 }
 
-def GPUExpandDimensionsPass :
-    InterfacePass<"iree-codegen-gpu-expand-dimensions", "mlir::FunctionOpInterface"> {
-  let summary = "Pass to expand tensor op dims based on `expand_dims` lowering_config";
-  let dependentDialects = [
-    "::mlir::iree_compiler::IREE::Util::UtilDialect"
-  ];
-}
-
 def GPUTensorTileToSerialLoopsPass :
     InterfacePass<"iree-codegen-gpu-tensor-tile-to-serial-loops", "mlir::FunctionOpInterface"> {
   let summary = "Pass to tile reduction dimensions for certain GPU ops";

compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td

@@ -36,7 +36,6 @@ iree_lit_test_suite(
             "gpu_distribute_forall.mlir",
             "gpu_distribute_scf_for.mlir",
             "gpu_distribute_shared_memory.mlir",
-            "gpu_expand_dimensions.mlir",
             "gpu_fuse_and_hoist_forall.mlir",
             "gpu_generalize_named_ops.mlir",
             "gpu_greedily_distribute_to_threads.mlir",

compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel

@@ -32,7 +32,6 @@ iree_lit_test_suite(
     "gpu_distribute_forall.mlir"
     "gpu_distribute_scf_for.mlir"
     "gpu_distribute_shared_memory.mlir"
-    "gpu_expand_dimensions.mlir"
     "gpu_fuse_and_hoist_forall.mlir"
     "gpu_generalize_named_ops.mlir"
     "gpu_greedily_distribute_to_threads.mlir"