[Mlir-commits] [mlir] [mlir][tosa] Fix slice fold when result type does not match input (PR #233515)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Luke Hutton (lhutton1)

<details>
<summary>Changes</summary>

When folding, the input should only be returned if the output type matches to prevent a crash.

Fixes https://github.com/llvm/llvm-project/issues/233055

---
Full diff: https://github.com/llvm/llvm-project/pull/233515.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp (+1-1) 
- (modified) mlir/test/Dialect/Tosa/canonicalize.mlir (+11) 


``````````diff
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index f97446c294247..340e53cc2c15c 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -2299,7 +2299,7 @@ OpFoldResult SliceOp::fold(FoldAdaptor adaptor) {
     }
 
     if (sizeMatchesInput)
-      return getInput1();
+      return foldToInputIfTypeMatches(getType(), getInput1());
   }
 
   // The following checks require the input to be a constant
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 48e429846b669..b102f202d07fd 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -874,6 +874,17 @@ func.func @slice_fold_static_dynamic(%arg0: tensor<?x4xf32>) -> tensor<?x4xf32>
 
 // -----
 
+// CHECK-LABEL: @slice_nofold_static_to_dynamic
+func.func @slice_nofold_static_to_dynamic(%arg0: tensor<3x4xf32>) -> tensor<?x4xf32> {
+  %start = tosa.const_shape values(dense<0> : tensor<2xindex>) : () -> !tosa.shape<2>
+  %size = tosa.const_shape values(dense<[3, 4]> : tensor<2xindex>) : () -> !tosa.shape<2>
+  // CHECK: tosa.slice
+  %slice = tosa.slice %arg0, %start, %size : (tensor<3x4xf32>, !tosa.shape<2>, !tosa.shape<2>) -> tensor<?x4xf32>
+  return %slice : tensor<?x4xf32>
+}
+
+// -----
+
 // CHECK-LABEL: @slice_nofold_static
 func.func @slice_nofold_static(%arg0: tensor<3x4xf32>) -> tensor<3x2xf32> {
   %0 = tosa.const_shape values(dense<[0, 0]> : tensor<2xindex>) : () -> !tosa.shape<2>

``````````

</details>


https://github.com/llvm/llvm-project/pull/233515