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

[llvmlistbot at llvm.org]

Author: Balaji V. Iyer
Date: 2024-03-25T16:05:09-05:00
New Revision: 5f1f9cfaa42a8dee59c9f3a354f973fd8cb003d7

URL: https://github.com/llvm/llvm-project/commit/5f1f9cfaa42a8dee59c9f3a354f973fd8cb003d7
DIFF: https://github.com/llvm/llvm-project/commit/5f1f9cfaa42a8dee59c9f3a354f973fd8cb003d7.diff

LOG: [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.

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
    mlir/test/Dialect/Vector/linearize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
index 7ca03537049812..38536de43f13f2 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 2cbf9bec7a4136..e865fcb7a419c5 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>
+}