Skip to content

Commit

Permalink
[mlir][Vector] Fix an assertion on failing cast in vector-transfer-fl…
Browse files Browse the repository at this point in the history
…atten-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.
  • Loading branch information
bviyer committed Mar 25, 2024
1 parent 0a42994 commit 5f1f9cf
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 2 deletions.
4 changes: 2 additions & 2 deletions mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
13 changes: 13 additions & 0 deletions mlir/test/Dialect/Vector/linearize.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -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>
}

0 comments on commit 5f1f9cf

Please sign in to comment.