[MLIR] Vector dialect: Address post-merge review comments on #111541

[bjacob]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Benoit Jacob (bjacob)

Changes

---

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

1 Files Affected:

• (modified) mlir/test/Dialect/Vector/vector-contiguous-extract-strided-slice-to-extract.mlir (+6-15)

diff --git a/mlir/test/Dialect/Vector/vector-contiguous-extract-strided-slice-to-extract.mlir b/mlir/test/Dialect/Vector/vector-contiguous-extract-strided-slice-to-extract.mlir
index 9147e7bf02581e..ea1bdedaf76286 100644
--- a/mlir/test/Dialect/Vector/vector-contiguous-extract-strided-slice-to-extract.mlir
+++ b/mlir/test/Dialect/Vector/vector-contiguous-extract-strided-slice-to-extract.mlir
@@ -1,34 +1,25 @@
 // RUN: mlir-opt --test-vector-contiguous-extract-strided-slice-to-extract %s | FileCheck %s
 
-// CHECK-LABEL: @extract_strided_slice_to_extract_i8
-// CHECK:       %[[EXTRACT:.+]] = vector.extract {{.*}}[0, 0, 0, 0] : vector<8xi8> from vector<8x1x1x2x8xi8>
-// CHECK:       return %[[EXTRACT]] :  vector<8xi8>
-func.func @extract_strided_slice_to_extract_i8(%arg0 : vector<8x1x1x2x8xi8>) -> vector<8xi8> {
-  %1 = vector.extract_strided_slice %arg0 {offsets = [0, 0, 0, 0, 0], sizes = [1, 1, 1, 1, 8], strides = [1, 1, 1, 1, 1]} : vector<8x1x1x2x8xi8> to vector<1x1x1x1x8xi8>
-  %2 = vector.shape_cast %1 : vector<1x1x1x1x8xi8> to vector<8xi8>
-  return %2 : vector<8xi8>
-}
-
-// CHECK-LABEL: @extract_strided_slice_to_extract_i32
+// CHECK-LABEL: @contiguous
 // CHECK:        %[[EXTRACT:.+]] = vector.extract {{.*}}[0, 0, 0, 0, 0] : vector<4xi32> from vector<8x1x2x1x1x4xi32>
 // CHECK:       return %[[EXTRACT]] :  vector<4xi32>
-func.func @extract_strided_slice_to_extract_i32(%arg0 : vector<8x1x2x1x1x4xi32>) -> vector<4xi32> {
+func.func @contiguous(%arg0 : vector<8x1x2x1x1x4xi32>) -> vector<4xi32> {
   %1 = vector.extract_strided_slice %arg0 {offsets = [0, 0, 0, 0, 0, 0], sizes = [1, 1, 1, 1, 1, 4], strides = [1, 1, 1, 1, 1, 1]} : vector<8x1x2x1x1x4xi32> to vector<1x1x1x1x1x4xi32>
   %2 = vector.shape_cast %1 : vector<1x1x1x1x1x4xi32> to vector<4xi32>
   return %2 : vector<4xi32>
 }
 
-// CHECK-LABEL: @extract_strided_slice_to_extract_i32_non_contiguous_1
+// CHECK-LABEL: @non_full_size
 // CHECK:        vector.extract_strided_slice
-func.func @extract_strided_slice_to_extract_i32_non_contiguous_1(%arg0 : vector<8x1x2x1x1x4xi32>) -> vector<2xi32> {
+func.func @non_full_size(%arg0 : vector<8x1x2x1x1x4xi32>) -> vector<2xi32> {
   %1 = vector.extract_strided_slice %arg0 {offsets = [0, 0, 0, 0, 0, 0], sizes = [1, 1, 1, 1, 1, 2], strides = [1, 1, 1, 1, 1, 1]} : vector<8x1x2x1x1x4xi32> to vector<1x1x1x1x1x2xi32>
   %2 = vector.shape_cast %1 : vector<1x1x1x1x1x2xi32> to vector<2xi32>
   return %2 : vector<2xi32>
 }
 
-// CHECK-LABEL: @extract_strided_slice_to_extract_i32_non_contiguous_2
+// CHECK-LABEL: @non_contiguous
 // CHECK:        vector.extract_strided_slice
-func.func @extract_strided_slice_to_extract_i32_non_contiguous_2(%arg0 : vector<8x1x2x1x1x4xi32>) -> vector<2xi32> {
+func.func @non_contiguous(%arg0 : vector<8x1x2x1x1x4xi32>) -> vector<2xi32> {
   %1 = vector.extract_strided_slice %arg0 {offsets = [0, 0, 0, 0, 0, 0], sizes = [1, 1, 2, 1, 1, 1], strides = [1, 1, 1, 1, 1, 1]} : vector<8x1x2x1x1x4xi32> to vector<1x1x2x1x1x1xi32>
   %2 = vector.shape_cast %1 : vector<1x1x2x1x1x1xi32> to vector<2xi32>
   return %2 : vector<2xi32>

[bjacob]

@banach-space , see responses on #111541 for the things not changed here, such as the shape_cast.

[banach-space]

That's true speed of light Benoit, thank you!

I've left two small suggestions here and under #111541 :

• use CHECK-NEXT, and
• skip checking strides altogether).

LGTM otherwise!

Apologies for being a bit pedantic, especially when it comes to tests. That's after spending quite some quality time with them 😅 These updates are much appreciated 🙏🏻

[bjacob]

Regarding the unit strides: Actually I was wrong earlier, as the pattern will match as soon as there is any full-size slice dimension. So I went ahead and added a strided testcase... but then it failed to validate, the validator requiring unit strides. And indeed, this is said in the doc: https://mlir.llvm.org/docs/Dialects/Vector/#vectorextract_strided_slice-vectorextractstridedsliceop :

At the moment strides must contain only 1s.
However, it sounds possible that strides might become flexible in the future, so to limit the risk of bugs popping up then, I'd like to keep the check, even if it can't be exercised now. That's also in line with other patterns in this file.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[banach-space]

Regarding the unit strides: Actually I was wrong earlier, as the pattern will match as soon as there is any full-size slice dimension. So I went ahead and added a strided testcase... but then it failed to validate, the validator requiring unit strides. And indeed, this is said in the doc: https://mlir.llvm.org/docs/Dialects/Vector/#vectorextract_strided_slice-vectorextractstridedsliceop :

At the moment strides must contain only 1s.
However, it sounds possible that strides might become flexible in the future, so to limit the risk of bugs popping up then, I'd like to keep the check, even if it can't be exercised now. That's also in line with other patterns in this file.

Yeah, so much for "strided" in extract_strided_slice 😂 Perhaps add a TODO somewhere ("Add a test for non-unit strides once these are supported")?

[banach-space]

Thanks!