[mlir][vector] Use notifyMatchFailure instead of assert in VectorLinearize #93590

akroviakov · 2024-05-28T17:55:46Z

As it was suggested, the assert is replaced by notifyMatchFailure for improved consistency.

llvmbot · 2024-05-28T17:56:14Z

@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Artem Kroviakov (akroviakov)

Changes

As it was suggested, the assert is replaced by notifyMatchFailure for improved consistency.

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

1 Files Affected:

(modified) mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp (+17-13)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
index 156bf742f6297..840fd384894df 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
@@ -152,9 +152,10 @@ struct LinearizeVectorExtractStridedSlice final
   matchAndRewrite(vector::ExtractStridedSliceOp extractOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type dstType = getTypeConverter()->convertType(extractOp.getType());
-    assert(!(extractOp.getVector().getType().isScalable() ||
-             cast<VectorType>(dstType).isScalable()) &&
-           "scalable vectors are not supported.");
+    if (extractOp.getVector().getType().isScalable() ||
+        cast<VectorType>(dstType).isScalable())
+      return rewriter.notifyMatchFailure(extractOp,
+                                         "scalable vectors are not supported.");
     if (!isLessThanTargetBitWidth(extractOp, targetVectorBitWidth))
       return rewriter.notifyMatchFailure(
           extractOp, "Can't flatten since targetBitWidth <= OpSize");
@@ -265,10 +266,11 @@ struct LinearizeVectorShuffle final
   matchAndRewrite(vector::ShuffleOp shuffleOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type dstType = getTypeConverter()->convertType(shuffleOp.getType());
-    assert(!(shuffleOp.getV1VectorType().isScalable() ||
-             shuffleOp.getV2VectorType().isScalable() ||
-             cast<VectorType>(dstType).isScalable()) &&
-           "scalable vectors are not supported.");
+    if (shuffleOp.getV1VectorType().isScalable() ||
+        shuffleOp.getV2VectorType().isScalable() ||
+        cast<VectorType>(dstType).isScalable())
+      return rewriter.notifyMatchFailure(shuffleOp,
+                                         "scalable vectors are not supported.");
     if (!isLessThanTargetBitWidth(shuffleOp, targetVectorBitWidth))
       return rewriter.notifyMatchFailure(
           shuffleOp, "Can't flatten since targetBitWidth <= OpSize");
@@ -336,9 +338,10 @@ struct LinearizeVectorExtract final
   matchAndRewrite(vector::ExtractOp extractOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type dstTy = getTypeConverter()->convertType(extractOp.getType());
-    assert(!(extractOp.getVector().getType().isScalable() ||
-             cast<VectorType>(dstTy).isScalable()) &&
-           "scalable vectors are not supported.");
+    if (extractOp.getVector().getType().isScalable() ||
+        cast<VectorType>(dstTy).isScalable())
+      return rewriter.notifyMatchFailure(extractOp,
+                                         "scalable vectors are not supported.");
     if (!isLessThanTargetBitWidth(extractOp, targetVectorBitWidth))
       return rewriter.notifyMatchFailure(
           extractOp, "Can't flatten since targetBitWidth <= OpSize");
@@ -395,9 +398,10 @@ struct LinearizeVectorInsert final
   matchAndRewrite(vector::InsertOp insertOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type dstTy = getTypeConverter()->convertType(insertOp.getDestVectorType());
-    assert(!(insertOp.getDestVectorType().isScalable() ||
-             cast<VectorType>(dstTy).isScalable()) &&
-           "scalable vectors are not supported.");
+    if (insertOp.getDestVectorType().isScalable() ||
+        cast<VectorType>(dstTy).isScalable())
+      return rewriter.notifyMatchFailure(insertOp,
+                                         "scalable vectors are not supported.");
 
     if (!isLessThanOrEqualTargetBitWidth(insertOp.getSourceType(),
                                          targetVectorBitWidth))

Garra1980 · 2024-05-28T19:00:35Z

cc @Hardcode84

banach-space

Thanks for addressing my suggestion so swiftly!

Could you also add some basic tests to show that these patterns do not trigger for scalable vectors?

banach-space · 2024-05-28T19:56:57Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

-    assert(!(shuffleOp.getV1VectorType().isScalable() ||
-             shuffleOp.getV2VectorType().isScalable() ||
-             cast<VectorType>(dstType).isScalable()) &&
-           "scalable vectors are not supported.");
+    if (shuffleOp.getV1VectorType().isScalable() ||


vector.shuffle does not support scalable vectors, so keeping an assert should be fine for this one:

https://mlir.llvm.org/docs/Dialects/Vector/#vectorshuffle-vectorshuffleop

A comment explaining the rationale for using assert rather than notifyMatchFailure would be welcome :)

+1 on test coverage.

Should be addressed now

hanhanW · 2024-05-30T23:27:17Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

@@ -152,9 +152,10 @@ struct LinearizeVectorExtractStridedSlice final
  matchAndRewrite(vector::ExtractStridedSliceOp extractOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Type dstType = getTypeConverter()->convertType(extractOp.getType());


Why don't we just cast dstType to ShapedType here? I think it carries more information/methods v.s. Type. And you don't need to cast it in the below if condition.

Could you please provide more detail on how ShapedType would not need a cast for cast<VectorType>(dstType).isScalable()? AFAIK ShapedType has no isScalable(), I do not see other places in the pattern where we could use ShapedType's information/methods that are not provided by Type.

Sorry for the confusion, what I meant is VectorType. So it could either be

VectorType dstType = getTypeConverter()->convertType(extractOp.getType());

or

auto dstType = cast<VectorType>(getTypeConverter()->convertType(extractOp.getType()));

you can actually do getTypeConverter()->convertType<VectorType>(...). Also, it's better to check convertType result for null.

Thanks for the clarification, should now be addressed

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

chencha3 · 2024-06-07T15:48:18Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

-             cast<VectorType>(dstTy).isScalable()) &&
-           "scalable vectors are not supported.");
+    if (extractOp.getVector().getType().isScalable() ||
+        cast<VectorType>(dstTy).isScalable())


Actually the dstTy here is not always a vector type. It could be a scalar type too.
e.g., vector.extract %1 [0, 0]: f32 from vector<1024x1024xf32>. So, cast(dstTy) may cause the crash.

Would you say the following is better?

if ( (auto vecDstTy = cast<VectorType>(dstTy) && vecDstTy.isScalable()) || extractOp.getVector().getType().isScalable() )

So, cast(dstTy) may cause the crash.

Good point, but let's stick to one change per PR 😅 My recommendation:

identify a test case that would indeed crash,

fix the crash and use the test from 1. for a follow-up PR.

WDYT?

What's the update here?

The test case is there:

%0 = vector.extract %arg0[1,1,1]: f32 from vector<2x8x2xf32>

but fixing it can get a bit tricky because right now there are a lot of vector result assumptions (e.g., isLessThanTargetBitWidth(), populateVectorLinearizeShuffleLikeOpsPatterns()), so yes, it should be another PR.
Any suggestions to make the fix least invasive are welcomed.

Any suggestions to make the fix least invasive are welcomed.

You should disable this pattern when the rank of the output is <= 1.

chencha3 · 2024-06-07T15:50:19Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

-           "scalable vectors are not supported.");
+    VectorType dstTy = getTypeConverter()->convertType<VectorType>(
+        insertOp.getDestVectorType());
+    assert(dstTy && "vector type destination expected.");


Same comments as above regarding to assert.

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

banach-space · 2024-06-09T17:47:56Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

-             cast<VectorType>(dstTy).isScalable()) &&
-           "scalable vectors are not supported.");
+    if (extractOp.getVector().getType().isScalable() ||
+        cast<VectorType>(dstTy).isScalable())


So, cast(dstTy) may cause the crash.

Good point, but let's stick to one change per PR 😅 My recommendation:

identify a test case that would indeed crash,

fix the crash and use the test from 1. for a follow-up PR.

WDYT?

banach-space · 2024-06-09T18:02:25Z

mlir/test/Dialect/Vector/linearize.mlir

-// ALL-SAME:     %[[VAL_0:.*]]: vector<[2]x[2]xf32>) -> vector<[2]x[2]xf32> {
-func.func @test_scalable_no_linearize(%arg0: vector<[2]x[2]xf32>) -> vector<[2]x[2]xf32> {
+// ALL-SAME:     %[[VAL_0:.*]]: vector<[2]x[2]xf32>,  %[[VAL_1:.*]]: vector<2x[2]xf32>) -> vector<[2]x[2]xf32> {
+func.func @test_scalable_no_linearize(%arg0: vector<[2]x[2]xf32>, %arg1: vector<2x[2]xf32>) -> vector<[2]x[2]xf32> {


I introduced this "function" to complement test_linearize:

llvm-project/mlir/test/Dialect/Vector/linearize.mlir

Lines 7 to 31 in e090bac

func.func @test_linearize(%arg0: vector<2x2xf32>) -> vector<2x2xf32> {

// DEFAULT: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x2xf32> to vector<4xf32>

// DEFAULT: %[[CST:.*]] = arith.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : vector<4xf32>

// DEFAULT: %[[RES:.*]] = vector.shape_cast %[[CST]] : vector<4xf32> to vector<2x2xf32>

// BW-128: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x2xf32> to vector<4xf32>

// BW-128: %[[CST:.*]] = arith.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : vector<4xf32>

// BW-128: %[[RES:.*]] = vector.shape_cast %[[CST]] : vector<4xf32> to vector<2x2xf32>

// BW-0: %[[RES:.*]] = arith.constant dense<{{.*}}> : vector<2x2xf32>

%0 = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : vector<2x2xf32>

// DEFAULT: %{{.*}} = math.sin %[[ARG]] : vector<4xf32>

// BW-128: %{{.*}} = math.sin %[[ARG]] : vector<4xf32>

// BW-0: %{{.*}} = math.sin %{{.*}} : vector<2x2xf32>

%1 = math.sin %arg0 : vector<2x2xf32>

// DEFAULT: %{{.*}} = arith.addf %[[ARG]], %[[CST]] : vector<4xf32>

// BW-128: %{{.*}} = arith.addf %[[ARG]], %[[CST]] : vector<4xf32>

// BW-0: %{{.*}} = arith.addf %{{.*}} : vector<2x2xf32>

%2 = arith.addf %arg0, %0 : vector<2x2xf32>

// ALL: return %[[RES]] : vector<2x2xf32>

return %0 : vector<2x2xf32>

}

. This way, for a set of tests for "fixed width" vectors it's quite easy to find to corresponding tests for "scalable vectors". I should've moved it next to test_linearize to make this clear - that's my bad, sorry for that!

With this in mind, would you be OK writing:

@test_extract_strided_slice_1_scalable,

`@test_extract_strided_slice_2_scalable,

and so on? The check lines could be as simple as:

// CHECK-LABEL: // CHECK-NOT: vector.shuffle // CHECK-NOT: vector.shape_cast // CHECK: vector.extract_strided_slice

(as in, the main thing to check would be that e.g. vector.shuffle Ops are not inserted).

I hope it is addressed now

The changes to the function signature should be reverted.

mlir/test/Dialect/Vector/linearize.mlir

banach-space · 2024-06-12T19:10:11Z

mlir/test/Dialect/Vector/linearize.mlir

-// ALL-SAME:     %[[VAL_0:.*]]: vector<[2]x[2]xf32>) -> vector<[2]x[2]xf32> {
-func.func @test_scalable_no_linearize(%arg0: vector<[2]x[2]xf32>) -> vector<[2]x[2]xf32> {
+// ALL-SAME:     %[[VAL_0:.*]]: vector<[2]x[2]xf32>,  %[[VAL_1:.*]]: vector<2x[2]xf32>) -> vector<[2]x[2]xf32> {
+func.func @test_scalable_no_linearize(%arg0: vector<[2]x[2]xf32>, %arg1: vector<2x[2]xf32>) -> vector<[2]x[2]xf32> {


The changes to the function signature should be reverted.

banach-space · 2024-06-12T19:11:49Z

mlir/test/Dialect/Vector/linearize.mlir

@@ -246,6 +257,16 @@ func.func @test_vector_extract(%arg0: vector<2x8x2xf32>) -> vector<8x2xf32> {
  return %0 : vector<8x2xf32>
 }

+// ALL-LABEL:   func.func @test_vector_extract_scalable(
+// ALL-SAME:    %[[VAL_0:.*]]: vector<2x[2]xf32>) -> f32 {
+func.func @test_vector_extract_scalable(%arg1: vector<2x[2]xf32>) -> f32 {


Why are types inside this tens and @test_vector_extract different? Is this in any way significant?

Changed the types for scalable tests to closely resemble normal ones.

mlir/test/Dialect/Vector/linearize.mlir

banach-space · 2024-06-13T10:49:34Z

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

-             cast<VectorType>(dstTy).isScalable()) &&
-           "scalable vectors are not supported.");
+    if (extractOp.getVector().getType().isScalable() ||
+        cast<VectorType>(dstTy).isScalable())


What's the update here?

akroviakov · 2024-06-19T13:43:12Z

Any further notes or can it be merged?

banach-space

LGTM, thank you for addressing my comments and for following up on this 🙏🏻

…arize

…arize (llvm#93590) As it was [suggested](llvm#92370 (comment)), the `assert` is replaced by `notifyMatchFailure` for improved consistency.

akroviakov requested review from hanhanW and nicolasvasilache as code owners May 28, 2024 17:55

llvmbot added mlir:vectorops mlir mlir:vector labels May 28, 2024

akroviakov mentioned this pull request May 28, 2024

[mlir][vector] Add support for linearizing Insert VectorOp in VectorLinearize #92370

Merged

Hardcode84 requested review from banach-space and Hardcode84 May 28, 2024 19:49

banach-space reviewed May 28, 2024

View reviewed changes

hanhanW reviewed May 30, 2024

View reviewed changes

akroviakov force-pushed the refactor_vector_linearize branch 2 times, most recently from 38689e2 to af59df9 Compare June 6, 2024 12:24

hanhanW approved these changes Jun 6, 2024

View reviewed changes

chencha3 reviewed Jun 7, 2024

View reviewed changes

banach-space reviewed Jun 9, 2024

View reviewed changes

akroviakov force-pushed the refactor_vector_linearize branch from af59df9 to b1b0384 Compare June 12, 2024 13:33

banach-space reviewed Jun 12, 2024

View reviewed changes

akroviakov force-pushed the refactor_vector_linearize branch from b1b0384 to ce9522c Compare June 13, 2024 09:26

banach-space reviewed Jun 13, 2024

View reviewed changes

akroviakov force-pushed the refactor_vector_linearize branch from ce9522c to 3392e43 Compare June 13, 2024 19:47

banach-space approved these changes Jun 19, 2024

View reviewed changes

[mlir][vector] Use notifyMatchFailure instead of assert in VectorLine…

71cfc85

…arize

akroviakov force-pushed the refactor_vector_linearize branch from 3392e43 to 71cfc85 Compare June 20, 2024 16:35

chencha3 merged commit 74a105a into llvm:main Jun 21, 2024
7 checks passed

	func.func @test_linearize(%arg0: vector<2x2xf32>) -> vector<2x2xf32> {
	// DEFAULT: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x2xf32> to vector<4xf32>
	// DEFAULT: %[[CST:.*]] = arith.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : vector<4xf32>
	// DEFAULT: %[[RES:.*]] = vector.shape_cast %[[CST]] : vector<4xf32> to vector<2x2xf32>

	// BW-128: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x2xf32> to vector<4xf32>
	// BW-128: %[[CST:.*]] = arith.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : vector<4xf32>
	// BW-128: %[[RES:.*]] = vector.shape_cast %[[CST]] : vector<4xf32> to vector<2x2xf32>

	// BW-0: %[[RES:.]] = arith.constant dense<{{.}}> : vector<2x2xf32>
	%0 = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : vector<2x2xf32>

	// DEFAULT: %{{.*}} = math.sin %[[ARG]] : vector<4xf32>
	// BW-128: %{{.*}} = math.sin %[[ARG]] : vector<4xf32>
	// BW-0: %{{.}} = math.sin %{{.}} : vector<2x2xf32>
	%1 = math.sin %arg0 : vector<2x2xf32>

	// DEFAULT: %{{.*}} = arith.addf %[[ARG]], %[[CST]] : vector<4xf32>
	// BW-128: %{{.*}} = arith.addf %[[ARG]], %[[CST]] : vector<4xf32>
	// BW-0: %{{.}} = arith.addf %{{.}} : vector<2x2xf32>
	%2 = arith.addf %arg0, %0 : vector<2x2xf32>

	// ALL: return %[[RES]] : vector<2x2xf32>
	return %0 : vector<2x2xf32>
	}

[mlir][vector] Use notifyMatchFailure instead of assert in VectorLinearize #93590

[mlir][vector] Use notifyMatchFailure instead of assert in VectorLinearize #93590

Uh oh!

Conversation

akroviakov commented May 28, 2024

Uh oh!

llvmbot commented May 28, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Garra1980 commented May 28, 2024

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

akroviakov Jun 13, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

akroviakov commented Jun 19, 2024

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

llvmbot commented May 28, 2024 •

edited

Loading

akroviakov Jun 13, 2024 •

edited

Loading