[mlir][Vector] Fix an assertion on failing cast in vector-transfer-flatten-patterns (#86030)

When the result is not a vectorType, there is an assert. This patch will
do the check and bail when the result is not a VectorType.

2 files changed, 15 insertions, 2 deletions

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
index 7ca0353..38536de 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
@@ -22,9 +22,9 @@ using namespace mlir;
 static bool isLessThanTargetBitWidth(Operation *op, unsigned targetBitWidth) {
   auto resultTypes = op->getResultTypes();
   for (auto resType : resultTypes) {
-    VectorType vecType = cast<VectorType>(resType);
+    VectorType vecType = dyn_cast<VectorType>(resType);
     // Reject index since getElementTypeBitWidth will abort for Index types.
-    if (vecType.getElementType().isIndex())
+    if (!vecType || vecType.getElementType().isIndex())
       return false;
     unsigned trailingVecDimBitWidth =
         vecType.getShape().back() * vecType.getElementTypeBitWidth();
diff --git a/mlir/test/Dialect/Vector/linearize.mlir b/mlir/test/Dialect/Vector/linearize.mlir
index 2cbf9be..e865fcb 100644
--- a/mlir/test/Dialect/Vector/linearize.mlir
+++ b/mlir/test/Dialect/Vector/linearize.mlir
@@ -90,3 +90,16 @@ func.func @test_index_no_linearize(%arg0: vector<2x2xindex>, %arg1: vector<2x2xi
     %0 = arith.addi %arg0, %arg1 : vector<2x2xindex>
     return %0 : vector<2x2xindex>
 }
+
+// -----
+
+// vectorizable operation (arith.mulf) with tensor result types.
+
+func.func @nonvec_result(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>) {
+    // CHECK: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
+    // CHECK128: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
+    // CHECK0: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
+    %0 = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
+
+    return %0, %arg0 : tensor<2x2xf32>, tensor<2x2xf32>
+}