Skip to content

[mlir][linalg] Refine test for tensor.pad vectorization #110742

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

Conversation

banach-space
Copy link
Contributor

@banach-space banach-space commented Oct 1, 2024

This particular test was missing the vectorize_padding attribute and
hence didn't really trigger anything meaningful. I've added the
attribute and additional check-lines to verify that:

  • The vectorizer was indeed run (and generated vector Ops).
  • The final tensor.insert_slice Op is preserved.

Note that another tensor.insert_slice Op, inserted by the vectoriser,
was rewritten as a vector.transfer_read/vector.transfer_write pair
and hence is not present in the test output.

This particular test was missing the `vectorize_padding` attribute and
hence didn't really trigger anything meaningful. I've added the
attribute and additional check-lines to verify that:
  * The vectorizer was indeed run (and generated vector Ops).
  * The final `tensor.insert_slice` Op is preserved. Note that another
    `tensor.insert_slice` Op, inserted by the vectoriser, was rewritten
    as a `vector.transfer_read/`vector.transfer_write` pair and hence is
    not present in the output..
@llvmbot
Copy link
Member

llvmbot commented Oct 1, 2024

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-linalg

Author: Andrzej Warzyński (banach-space)

Changes

This particular test was missing the vectorize_padding attribute and
hence didn't really trigger anything meaningful. I've added the
attribute and additional check-lines to verify that:

  • The vectorizer was indeed run (and generated vector Ops).
  • The final tensor.insert_slice Op is preserved. Note that another
    tensor.insert_slice Op, inserted by the vectoriser, was rewritten
    as a vector.transfer_read/vector.transfer_write` pair and hence is
    not present in the output..

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

1 Files Affected:

  • (modified) mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir (+8-2)
diff --git a/mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir b/mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir
index 9a43d43cd9460a..80564ad35cfdbe 100644
--- a/mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir
@@ -1091,9 +1091,15 @@ module attributes {transform.with_named_sequence} {
 func.func private @make_vector() -> tensor<12x13xf32>
 
 // CHECK-LABEL: func @pad_and_insert_slice_dest
+//  CHECK-SAME:     %[[ARG0:.*]]: tensor<1x5x6xf32>
 // Check the insert slice is not rewritten if the padded result is used by the destination operand.
+//   CHECK-NOT:   tensor.pad
+//       CHECK:   %[[EMPTY:.*]] = tensor.empty() : tensor<1x12x13xf32>
+//       CHECK:   %[[WRITE_1:.*]] = vector.transfer_write %{{.*}}, %[[EMPTY]]{{.*}} : vector<1x12x13xf32>, tensor<1x12x13xf32>
+//       CHECK:   %[[READ:.*]]  = vector.transfer_read %[[ARG0:.*]]{{.*}} : tensor<1x5x6xf32>, vector<1x5x6xf32>
+//       CHECK:   %[[WRITE_2:.*]] = vector.transfer_write %[[READ]], %[[WRITE_1]]{{.*}} : vector<1x5x6xf32>, tensor<1x12x13xf32>
 //       CHECK:   %[[T1:.*]] = call @make_vector() : () -> tensor<12x13xf32>
-//       CHECK:   = tensor.insert_slice %[[T1]] into
+//       CHECK:   tensor.insert_slice %[[T1]] into %[[WRITE_2]]
 func.func @pad_and_insert_slice_dest(
     %arg0: tensor<1x5x6xf32>) -> tensor<1x12x13xf32> {
   %c5 = arith.constant 5.0 : f32
@@ -1110,7 +1116,7 @@ module attributes {transform.with_named_sequence} {
   transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
     %3 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op
     %4 = transform.get_parent_op %3 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
-    %5 = transform.structured.vectorize_children_and_apply_patterns %4 : (!transform.any_op) -> !transform.any_op
+    %5 = transform.structured.vectorize_children_and_apply_patterns %4  { vectorize_padding } : (!transform.any_op) -> !transform.any_op
     transform.yield
   }
 }

Copy link
Contributor

@hanhanW hanhanW left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks

@banach-space banach-space merged commit 33c134e into llvm:main Oct 2, 2024
11 checks passed
@banach-space banach-space deleted the andrzej/update_test_tensor_pad branch October 2, 2024 06:40
Sterling-Augustine pushed a commit to Sterling-Augustine/llvm-project that referenced this pull request Oct 3, 2024
This particular test was missing the `vectorize_padding` attribute and
hence didn't really trigger anything meaningful. I've added the
attribute and additional check-lines to verify that:
  * The vectorizer was indeed run (and generated vector Ops).
  * The final `tensor.insert_slice` Op is preserved.

Note that another `tensor.insert_slice` Op, inserted by the vectoriser,
was rewritten as a `vector.transfer_read`/`vector.transfer_write` pair
and hence is not present in the test output.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants