[Mlir-commits] [mlir] Revert "[tosa]: canonicalize dynamic size of tosa.slice to static output shape" (PR #135525)

[Thurston Dang]

https://github.com/thurstond created https://github.com/llvm/llvm-project/pull/135525

Reverts llvm/llvm-project#135429 due buildbot breakage: https://lab.llvm.org/buildbot/#/builders/169/builds/10405

Based on the ASan output, I think after the replaceOp on line 775, it's no longer valid to do getSize() on sliceOp:
```
   775      rewriter.replaceOp(sliceOp, newSliceOp.getResult());
   776
   777      // Remove const_shape size op when it no longer has use point.
   778      Operation *sizeConstShape = sliceOp.getSize().getDefiningOp();
```


>From 61393cc94c0223c2d2fe3ee15e6f79e55c4f3e8b Mon Sep 17 00:00:00 2001
From: Thurston Dang <thurston.dang at gmail.com>
Date: Sun, 13 Apr 2025 00:06:12 -0700
Subject: [PATCH] =?UTF-8?q?Revert=20"[tosa]:=20canonicalize=20dynamic=20si?=
 =?UTF-8?q?ze=20of=20tosa.slice=20to=20static=20output=20shap=E2=80=A6"?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This reverts commit 60b1d44d708d5912897c7bdab681cbefe8f35e79.
---
 .../Dialect/Tosa/IR/TosaCanonicalizations.cpp | 55 +------------------
 mlir/test/Dialect/Tosa/canonicalize.mlir      | 15 -----
 2 files changed, 1 insertion(+), 69 deletions(-)

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 84f89bfd7f2d3..c4ef7d0bb9ff5 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -731,62 +731,9 @@ struct ConcatSliceOptimization : public OpRewritePattern<tosa::SliceOp> {
   }
 };
 
-// Update size operand of tosa.slice if size has dynamic dims but corresponding
-// output dim is static
-struct SliceDynamicSizeCanonicalization
-    : public OpRewritePattern<tosa::SliceOp> {
-  using OpRewritePattern<tosa::SliceOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(tosa::SliceOp sliceOp,
-                                PatternRewriter &rewriter) const override {
-    ShapedType resultType = cast<ShapedType>(sliceOp.getType());
-
-    ElementsAttr sizeElems;
-    if (!matchPattern(sliceOp.getSize(), m_Constant(&sizeElems))) {
-      return rewriter.notifyMatchFailure(
-          sliceOp, "size of slice must be a static ranked shape");
-    }
-
-    llvm::SmallVector<int64_t> sliceSizes =
-        llvm::to_vector(sizeElems.getValues<int64_t>());
-
-    bool replaceSliceSize{false};
-    // if size op has -1 indicating dynamic shape but corresponding dim on the
-    // output is statically known, update size to match with known output dim
-    // shape
-    for (const auto &[index, size] : llvm::enumerate(sliceSizes)) {
-      if (size == -1 && !resultType.isDynamicDim(index)) {
-        sliceSizes[index] = resultType.getDimSize(index);
-        replaceSliceSize = true;
-      }
-    }
-
-    if (!replaceSliceSize) {
-      return rewriter.notifyMatchFailure(
-          sliceOp, "no dimension of size of slice is dynamic that resolves "
-                   "to static output shape");
-    }
-
-    auto size_op = getTosaConstShape(rewriter, sliceOp.getLoc(), sliceSizes);
-    auto newSliceOp = rewriter.create<tosa::SliceOp>(
-        sliceOp.getLoc(), sliceOp.getType(), sliceOp.getInput1(),
-        sliceOp.getStart(), size_op);
-
-    rewriter.replaceOp(sliceOp, newSliceOp.getResult());
-
-    // Remove const_shape size op when it no longer has use point.
-    Operation *sizeConstShape = sliceOp.getSize().getDefiningOp();
-    if (sizeConstShape->getResult(0).hasOneUse())
-      rewriter.eraseOp(sizeConstShape);
-
-    return success();
-  }
-};
-
 void SliceOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {
-  results.add<ConcatSliceOptimization, SliceDynamicSizeCanonicalization>(
-      context);
+  results.add<ConcatSliceOptimization>(context);
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index a754a46be603f..b366b4f1e4fd4 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1212,18 +1212,3 @@ func.func @do_not_fold_intdiv_division_by_0() -> tensor<1x24x2xi32> {
   %16 = tosa.intdiv %4, %1 : (tensor<1x24x2xi32>, tensor<1x24x2xi32>) -> tensor<1x24x2xi32>
   return %16 : tensor<1x24x2xi32>
 }
-
-
-// ----
-// CHECK-LABEL:   func.func @slice_dynamic_size_static_output_canonicalize(
-// CHECK-SAME:                     %[[ARG0:.*]]: tensor<2x60x59x?xf32>) -> tensor<2x60x58x?xf32> {
-// CHECK:           %[[START:.*]] = tosa.const_shape  {values = dense<0> : tensor<4xindex>} : () -> !tosa.shape<4>
-// CHECK:           %[[SIZE:.*]] = tosa.const_shape  {values = dense<[2, 60, 58, -1]> : tensor<4xindex>} : () -> !tosa.shape<4>
-// CHECK:           %[[SLICE:.*]] = tosa.slice %[[ARG0]], %[[START]], %[[SIZE]] : (tensor<2x60x59x?xf32>, !tosa.shape<4>, !tosa.shape<4>) -> tensor<2x60x58x?xf32>
-// CHECK:           return %[[SLICE]]
-func.func @slice_dynamic_size_static_output_canonicalize(%arg0: tensor<2x60x59x?xf32>) -> tensor<2x60x58x?xf32> {
-    %0 = tosa.const_shape  {values = dense<0> : tensor<4xindex>} : () -> !tosa.shape<4>
-    %1 = tosa.const_shape  {values = dense<[-1, 60, 58, -1]> : tensor<4xindex>} : () -> !tosa.shape<4>
-    %2 = tosa.slice %arg0, %0, %1 : (tensor<2x60x59x?xf32>, !tosa.shape<4>, !tosa.shape<4>) -> tensor<2x60x58x?xf32>
-    return %2 : tensor<2x60x58x?xf32>
-  }