From af1768b649ceadd35439a1abbf38f700e3132f60 Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Thu, 9 Jan 2025 18:43:57 +0000 Subject: [PATCH 1/2] [mlir][linalg] Add vectorization to the e2e test for `tensor.unpack` Following on from #122927 + #123031 that added support for masked vectorization of `tensor.insert_slice`, this PR extends the e2e test for `tensor.unpack` to leverage the new functionality. --- .../Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir | 4 ++-- .../Linalg/CPU/unpack-dynamic-inner-tile.mlir | 12 +++++++++--- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir index 15edae8b6d3f8..9a8708067c48f 100644 --- a/mlir/test/Integration/Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir +++ b/mlir/test/Integration/Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir @@ -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 %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 diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir index 4395dfe74914e..90ecf4abfd1d9 100644 --- a/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir +++ b/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir @@ -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 @@ -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. 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 + // 3. Bufferize before lowering to LLVM %bufferize = transform.bufferization.one_shot_bufferize %module {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op - // 4. Canonicalize + // 4. 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 From eab99bb6bd912ef1c20db78ec7d23c24ab8f2d3f Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Wed, 5 Feb 2025 17:43:32 +0000 Subject: [PATCH 2/2] fixup! [mlir][linalg] Add vectorization to the e2e test for `tensor.unpack` Fix numbering --- .../Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir index 90ecf4abfd1d9..a97238e5eeea4 100644 --- a/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir +++ b/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir @@ -99,11 +99,11 @@ module @transforms attributes { transform.with_named_sequence } { %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 - // 3. Bufferize before lowering to LLVM + // 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