From 6d6e67a294f7b670b13b5ae2ffece71fc79d576f Mon Sep 17 00:00:00 2001
From: "liqin.weng" <liqin.weng@spacemit.com>
Date: Tue, 10 Oct 2023 09:44:59 +0800
Subject: [PATCH] [MLIR][TOSA] Add tosa.slice operation conversion failure
 scenario

Fixes #68481, In the following scenario, the conversion fails:
1. resultType of tosa.slice is UnrankedTensorType
2. tosa.slice.getsize().size() < resultType.getRank()
---
 mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp |  4 ++++
 .../TosaToTensor/tosa-to-tensor-invalid.mlir      | 15 +++++++++++++++
 2 files changed, 19 insertions(+)
 create mode 100644 mlir/test/Conversion/TosaToTensor/tosa-to-tensor-invalid.mlir

diff --git a/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp b/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
index f51ada8d08b5e..7a62429d37004 100644
--- a/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
+++ b/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
@@ -243,6 +243,10 @@ class SliceConverter : public OpConversionPattern<tosa::SliceOp> {
                   ConversionPatternRewriter &rewriter) const final {
     Location loc = sliceOp.getLoc();
     Value input = adaptor.getInput();
+    ShapedType resultType = cast<ShapedType>(sliceOp.getType());
+    if (llvm::isa<UnrankedTensorType>(resultType) ||
+        resultType.getRank() != static_cast<int64_t>(sliceOp.getSize().size()))
+      return failure();
     SmallVector<int64_t> strides, sizes;
     ArrayRef<int64_t> starts = sliceOp.getStart();
     strides.resize(cast<ShapedType>(sliceOp.getType()).getRank(), 1);
diff --git a/mlir/test/Conversion/TosaToTensor/tosa-to-tensor-invalid.mlir b/mlir/test/Conversion/TosaToTensor/tosa-to-tensor-invalid.mlir
new file mode 100644
index 0000000000000..cfa1c41f24679
--- /dev/null
+++ b/mlir/test/Conversion/TosaToTensor/tosa-to-tensor-invalid.mlir
@@ -0,0 +1,15 @@
+// RUN: mlir-opt --split-input-file -pass-pipeline="builtin.module(func.func(tosa-to-tensor))" %s -verify-diagnostics
+
+// CHECK-LABEL:  @slice_resultType_unranked
+func.func @slice_resultType_unranked(%arg0: tensor<?xf32>) -> (tensor<*xf32>) {
+  // expected-error@+1 {{failed to legalize operation 'tosa.slice'}}
+  %0 = "tosa.slice"(%arg0) {start = array<i64: 2>, size = array<i64: 0>} : (tensor<?xf32>)  -> (tensor<*xf32>)
+  return %0 : tensor<*xf32>
+}
+
+// CHECK-LABEL:  @slice_resultRank_neq_opSize
+func.func @slice_resultRank_neq_opSize(%arg0: tensor<12xf32>) -> (tensor<2xf32>) {
+  // expected-error@+1 {{failed to legalize operation 'tosa.slice'}}
+  %0 = "tosa.slice"(%arg0) {start = array<i64: 2>, size = array<i64: 2, 3>} : (tensor<12xf32>)  -> (tensor<2xf32>)
+  return %0 : tensor<2xf32>
+}