Skip to content

[mlir][linalg] Add vectorization to the e2e test for tensor.unpack #123032

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.

Sign up for GitHub

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
merged 2 commits into from
Feb 28, 2025
Merged

[mlir][linalg] Add vectorization to the e2e test for tensor.unpack #123032

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
af1768b
[mlir][linalg] Add vectorization to the e2e test for `tensor.unpack`
banach-space Jan 9, 2025
eab99bb
fixup! [mlir][linalg] Add vectorization to the e2e test for `tensor.u…
banach-space Feb 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions mlir/test/Integration/Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,11 +121,11 @@ module @transforms attributes { transform.with_named_sequence } {
transform.apply_patterns.canonicalization
} : !transform.op<"func.func">

// 3. Bufferize before lowering to LLVM
// 4. Bufferize before lowering to LLVM
Copy link
Contributor

Choose a reason for hiding this comment

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

This is why I prefer not using numbers. They always end up wrong when code is moved around or when code is inserted/removed. If it were me, I'd remove all the numbers.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

They always end up wrong when code is moved around or when code is inserted/removed.

Blame the PR author 😅

This is a very good point and, in principle, I agree. To me, the numbering itself is secondary. However, these numbers provide clear separation of high-level steps - that's far more valuable to me and that's what I want to document here. Is there a better way than numbers?

    // Use HS to mark High-Level STEP
    
    // HS: Tile so that we can decompose tensor.pack
    // HS: Decompose the tiled unpack Op into tensor.extract_slice + tensor.insert_slice
    // HS: Vectorize tensor.insert_slice
    // HS: Bufferize before lowering to LLVM
    // HS: Canonicalize
    
    // Use # to mark High-Level STEP
    
    // # Tile so that we can decompose tensor.pack
    // # Decompose the tiled unpack Op into tensor.extract_slice + tensor.insert_slice
    // # Vectorize tensor.insert_slice
    // # Bufferize before lowering to LLVM
    // # Canonicalize

    // Use caps to mark High-Level STEP
    
    // TILE (so that we can decompose tensor.pack)
    // DECOMPOSE THE TILED UNPACK OP (into tensor.extract_slice + tensor.insert_slice)
    // VECTORIZE
    // BUFFERIZE
    // CANONICALIZE

It won't matter that much in the context in this single test, but perhaps it's something to consider in general. Having said that, "numbers" is what we tend to use in MLIR for marking high-level steps 🤷🏻

%bufferize = transform.bufferization.one_shot_bufferize %module
{bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op

// 4. Canonicalize
// 5. Canonicalize
%func_op_bufferized = transform.structured.match ops{["func.func"]} in %bufferize : (!transform.any_op) -> !transform.op<"func.func">
transform.apply_patterns to %func_op_bufferized {
transform.apply_patterns.canonicalization
Expand Down
14 changes: 10 additions & 4 deletions mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule |\
// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule \
// DEFINE: --lower-vector-mask |\
// DEFINE: mlir-opt \
// DEFINE: -test-lower-to-llvm -o %t
// DEFINE: %{entry_point} = main
Expand Down Expand Up @@ -90,14 +91,19 @@ module @transforms attributes { transform.with_named_sequence } {
%func_op = transform.get_parent_op %tiled_pack_op_p {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">
transform.apply_patterns to %func_op {
transform.apply_patterns.linalg.decompose_pack_unpack
transform.apply_patterns.linalg.decompose_pad
transform.apply_patterns.canonicalization
} : !transform.op<"func.func">

// 3. Bufferize before lowering to LLVM
// 3. Vectorize tensor.insert_slice
// Vector sizes match the inner tiles in the payload IR.
%slice = transform.structured.match ops{["tensor.insert_slice"]} in %func_op : (!transform.op<"func.func">) -> !transform.any_op
transform.structured.vectorize %slice vector_sizes [8, 1] : !transform.any_op

// 4. Bufferize before lowering to LLVM
%bufferize = transform.bufferization.one_shot_bufferize %module
{bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op

// 4. Canonicalize
// 5. Canonicalize
%func_op_bufferized = transform.structured.match ops{["func.func"]} in %bufferize : (!transform.any_op) -> !transform.op<"func.func">
transform.apply_patterns to %func_op_bufferized {
transform.apply_patterns.canonicalization
Expand Down