Revert (2nd) of "Propagate reshapes through generics with reduction" (#19647)
This reverts commit a5c3879bdd0779de2858f8095771dcc7e353b544.
Seems like there is a regression due to dynamic shapes
https://github.com/iree-org/iree/pull/18968#issuecomment-2579316594
Signed-off-by: MaheshRavishankar <mravisha@amd.com>
diff --git a/.github/workflows/pkgci_regression_test.yml b/.github/workflows/pkgci_regression_test.yml
index d194bba..448c556 100644
--- a/.github/workflows/pkgci_regression_test.yml
+++ b/.github/workflows/pkgci_regression_test.yml
@@ -125,7 +125,7 @@
--goldentime-rocm-vae-ms 310.0 \
--goldendispatch-rocm-unet 1602 \
--goldendispatch-rocm-clip 1139 \
- --goldendispatch-rocm-vae 245 \
+ --goldendispatch-rocm-vae 246 \
--goldensize-rocm-unet-bytes 2280000 \
--goldensize-rocm-clip-bytes 860000 \
--goldensize-rocm-vae-bytes 840000 \
@@ -150,7 +150,7 @@
--goldentime-rocm-vae-ms 75.0 \
--goldendispatch-rocm-unet 1602 \
--goldendispatch-rocm-clip 1139 \
- --goldendispatch-rocm-vae 245 \
+ --goldendispatch-rocm-vae 246 \
--goldensize-rocm-unet-bytes 2270000 \
--goldensize-rocm-clip-bytes 860000 \
--goldensize-rocm-vae-bytes 840000 \
diff --git a/compiler/src/iree/compiler/DispatchCreation/BubbleUpExpandShapes.cpp b/compiler/src/iree/compiler/DispatchCreation/BubbleUpExpandShapes.cpp
index 1f8e010..71fe957 100644
--- a/compiler/src/iree/compiler/DispatchCreation/BubbleUpExpandShapes.cpp
+++ b/compiler/src/iree/compiler/DispatchCreation/BubbleUpExpandShapes.cpp
@@ -134,8 +134,12 @@
return false;
}
+ // Do not fuse producer generic op if it has more than one user
+ // or any reduction iterators.
if (auto producerGenericOp = dyn_cast<linalg::GenericOp>(producer)) {
- return true;
+ return producerGenericOp->hasOneUse() &&
+ llvm::all_of(producerGenericOp.getIteratorTypesArray(),
+ linalg::isParallelIterator);
}
// Do not fuse with any producer linalg named ops for now.
@@ -143,9 +147,11 @@
return false;
}
- // Do not fuse with consumer linalg named ops.
+ // Do not fuse with consumer linalg named ops or reductions.
if (auto consumerLinalgOp = dyn_cast<linalg::LinalgOp>(consumer)) {
- return isa<linalg::GenericOp>(consumerLinalgOp);
+ return isa<linalg::GenericOp>(consumerLinalgOp) &&
+ llvm::all_of(consumerLinalgOp.getIteratorTypesArray(),
+ linalg::isParallelIterator);
}
// Fuse in all other cases.
return true;
diff --git a/compiler/src/iree/compiler/DispatchCreation/BubbleUpExtractSlices.cpp b/compiler/src/iree/compiler/DispatchCreation/BubbleUpExtractSlices.cpp
index 672ef9a..47d1699 100644
--- a/compiler/src/iree/compiler/DispatchCreation/BubbleUpExtractSlices.cpp
+++ b/compiler/src/iree/compiler/DispatchCreation/BubbleUpExtractSlices.cpp
@@ -149,7 +149,6 @@
patterns.insert<BubbleUpExtract>(context);
patterns.insert<SwapExtractSliceOfFill>(context);
tensor::populateFoldTensorEmptyPatterns(patterns, false);
- linalg::FillOp::getCanonicalizationPatterns(patterns, context);
if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
return signalPassFailure();
}
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/bubble_up_extract_slice.mlir b/compiler/src/iree/compiler/DispatchCreation/test/bubble_up_extract_slice.mlir
index c5311c2..b582b56 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/bubble_up_extract_slice.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/bubble_up_extract_slice.mlir
@@ -141,23 +141,3 @@
// CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]] :
// CHECK-SAME: outs(%[[EMPTY]] :
// CHECK: return %[[GENERIC]]
-
-// -----
-
-func.func @fold_extract_of_expand_of_fill(%arg0 : index, %arg1 : index, %arg2 : index) -> tensor<?xf16> {
- %cst0 = arith.constant 0.0 : f16
- %0 = tensor.empty(%arg0) : tensor<?xf16>
- %2 = linalg.fill ins(%cst0 : f16) outs(%0 : tensor<?xf16>) -> tensor<?xf16>
- %3 = tensor.expand_shape %2 [[0, 1]] output_shape[1, %arg1] : tensor<?xf16> into tensor<1x?xf16>
- %4 = tensor.extract_slice %3 [0, 0] [1, %arg2] [1, 1] : tensor<1x?xf16> to tensor<?xf16>
- func.return %4 : tensor<?xf16>
-}
-
-// CHECK-LABEL: func @fold_extract_of_expand_of_fill
-// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index
-// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index
-// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index
-// CHECK-DAG: %[[EMPTY:.+]] = tensor.empty(%[[ARG2]])
-// CHECK-DAG: %[[CST0:.+]] = arith.constant 0.0
-// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST0]] : f16) outs(%[[EMPTY]]
-// CHECK: return %[[FILL]]
