dont clone extract_slice

Signed-off-by: Ian Wood <[email protected]>
iree-org · May 30, 2024 · 47c41fd · 47c41fd
1 parent aa0bc40
commit 47c41fd
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp
@@ -601,9 +601,15 @@ bool isClonableIntoDispatchOp(Operation *op) {
   // with bufferization. Make them clonable when fixed.
   if (isa<affine::AffineApplyOp, arith::IndexCastOp, linalg::FillOp,
           tensor::EmptyOp, tensor::CastOp, tensor::ExtractOp,
-          tensor::ExtractSliceOp, complex::CreateOp>(op)) {
+          complex::CreateOp>(op)) {
     return true;
   }
+
+  if (auto extractSlice = dyn_cast<tensor::ExtractSliceOp>(op)) {
+    ArrayRef<int64_t> shape = extractSlice.getResult().getType().getShape();
+    return shape.size() == 0 || shape[0] < 10;
+  }
+
   if (isDequantizationLikeOp(op)) {
     return true;
   }