<table border="1" cellspacing="0" cellpadding="8">
<tr>
<th>Issue</th>
<td>
<a href=https://github.com/llvm/llvm-project/issues/145861>145861</a>
</td>
</tr>
<tr>
<th>Summary</th>
<td>
[MLIR][linalg] `GeneralizeOuterUnitDimsPackOpPattern` not checking trailing dimension for tiling
</td>
</tr>
<tr>
<th>Labels</th>
<td>
mlir
</td>
</tr>
<tr>
<th>Assignees</th>
<td>
</td>
</tr>
<tr>
<th>Reporter</th>
<td>
rYm-A
</td>
</tr>
</table>
<pre>
Compiling an ML model with IREE I found that an incorrect linalg.transpose op is generated after `iree-codegen-decompose-pack-unpack-ops`.
The linalg.pack I'm attempting to lower before `iree-odegen-decompose-pack-unpack-ops` is:
```
%pack = linalg.pack %extracted_slice_1
outer_dims_perm = [1, 0, 2]
inner_dims_pos = [0, 2]
inner_tiles = [8, 1]
into %extracted_slice_2
{lowering_config = #iree_codegen.lowering_config<
tile_sizes = [[1, 48, 64], [1, 1, 1]]>}
: tensor<8x1x1xf32> -> tensor<1x1x1x8x1xf32>
```
And the result is:
```
%18 = "tensor.empty"() : () -> tensor<1x1x1xf32>
%19 = "linalg.transpose"(%16, %18) <{permutation = array<i64: 0, 0, 2>}> ({
^bb0(%arg7: f32, %arg8: f32):
"linalg.yield"(%arg7) : (f32) -> ()
}) : (tensor<8x1x1xf32>, tensor<1x1x1xf32>) -> tensor<1x1x1xf32>
```
Note that the permutation array is incorrect, what generates an assertion error:
` Assertion permutationMap.isPermutation() && "Invalid permutation vector"' failed.`
Checking this [PR](https://github.com/llvm/llvm-project/pull/115312) by @banach-space I found that this pattern shouldn't be matched by GeneralizeOuterUnitDimsPackOpPattern, since the `inner_dim_pos` of this linalg.pack op aren't the last of the source, which was assumed for this PR.
The check in https://github.com/llvm/llvm-project/blame/237b8de2c0d9ee50c6a744e95c0706c8cdea70e1/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp#L1183 seems not to catch this particular example since the `-1` at the end may be accepting the N + 1 trailing dims, and not the N trailing dims. Could you confirm if this is intended?
Changing the check to `return dimPos >= (srcRank - numTiles);` will cause the pattern in this particular example to fail, and as a result, this linalg.op will eventually be lowered to:
```
%expanded = tensor.expand_shape %arg0 [[0, 1], [2], [3, 4]] output_shape [1, 8, 1, 1, 1] : tensor<8x1x1xf32> into tensor<1x8x1x1x1xf32>
%transposed = linalg.transpose ins(%expanded : tensor<1x8x1x1x1xf32>) outs(%arg1 : tensor<1x1x1x8x1xf32>) permutation = [2, 0, 3, 1, 4] {lowering_config = #config}
```
### How to reproduce
IREE v3.5.0.
```
iree-opt packOp.mlir \
--pass-pipeline="builtin.module(func.func(iree-codegen-decompose-pack-unpack-ops))" \
--debug \
--mlir-disable-threading
```
Trace:
```
Stack dump without symbol names (ensure you have llvm-symbolizer in your PATH or set the environment var `LLVM_SYMBOLIZER_PATH` to point to it):
0 libIREECompiler.so 0x00007fec4b0a85b8 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) + 40
1 libIREECompiler.so 0x00007fec4b0a635e llvm::sys::RunSignalHandlers() + 238
2 libIREECompiler.so 0x00007fec4b0a8cb6
3 libc.so.6 0x00007fec454e6520
4 libIREECompiler.so 0x00007fec4b0d50b3 mlir::AffineMap::getContext() const + 3
5 libIREECompiler.so 0x00007fec4b102bc4 mlir::AffineMapAttr::get(mlir::AffineMap) + 20
6 libIREECompiler.so 0x00007fec4b0ff979 mlir::Builder::getAffineMapArrayAttr(llvm::ArrayRef<mlir::AffineMap>) + 137
7 libIREECompiler.so 0x00007fec4f9b71fb mlir::linalg::TransposeOp::getIndexingMaps() + 475
8 libIREECompiler.so 0x00007fec4cb25b76
9 libIREECompiler.so 0x00007fec4fc51f67
10 libIREECompiler.so 0x00007fec4f99da41 mlir::linalg::LinalgOp::getIndexingMapsArray() + 17
11 libIREECompiler.so 0x00007fec4fb1cb0d
12 libIREECompiler.so 0x00007fec4ff8e3cd
13 libIREECompiler.so 0x00007fec4ff8b664 mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<llvm::LogicalResult (mlir::Pattern const&)>) + 820
14 libIREECompiler.so 0x00007fec4ff729e7
15 libIREECompiler.so 0x00007fec4ff707bb mlir::applyPatternsGreedily(mlir::Region&, mlir::FrozenRewritePatternSet const&, mlir::GreedyRewriteConfig, bool*) + 1819
16 libIREECompiler.so 0x00007fec4dc12c0b
17 libIREECompiler.so 0x00007fec4dc1195c
18 libIREECompiler.so 0x00007fec4b31e5cb mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) + 635
19 libIREECompiler.so 0x00007fec4b31f1b9 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) + 329
20 libIREECompiler.so 0x00007fec4b321907 mlir::detail::OpToOpPassAdaptor::runOnOperationImpl(bool) + 2311
21 libIREECompiler.so 0x00007fec4b31ea8c mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) + 1852
22 libIREECompiler.so 0x00007fec4b3224db mlir::PassManager::run(mlir::Operation*) + 1531
23 libIREECompiler.so 0x00007fec4b3150be
24 libIREECompiler.so 0x00007fec4b314d3d
25 libIREECompiler.so 0x00007fec4b316c92 mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<llvm::LogicalResult (std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>, llvm::raw_ostream&, llvm::StringRef, llvm::StringRef) + 818
26 libIREECompiler.so 0x00007fec4b30f362 mlir::MlirOptMain(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&) + 226
27 libIREECompiler.so 0x00007fec4b008e1d ireeOptRunMain + 2013
28 libc.so.6 0x00007fec454cdd90
29 libc.so.6 0x00007fec454cde40 __libc_start_main + 128
30 iree-opt 0x000000000020169e
```
</pre>
<img width="1" height="1" alt="" src="http://email.email.llvm.org/o/eJzMWVtv4zpw_jXMC2GDoixZesiD4yTnBEhOguy2wOmLQVFjm12JVEkqiffXF0NdLOdyvNsWbQ3DscXh3DnfDCOcUzsNcEmSK5JcX4jW7429tH_Xs9VFYcrD5drUjaqU3lGh6cM9rU0JFX1Vfk_vnm9u6B3dmlaX1O-FRxKlpbEWpKeV0qLazb0V2jXGATUNVY7uQIMVHkoqth4sJSlTFmAmTQk70LMSpKmRftYI-WPW6vDHNI6kbE7YirDV9z0M3HGR3hG-rKnwHurGo6re0Mq8gqUFbI2FUcYviKDKkXjVySEp699sRXgSZJH4-kQ24Qm8eSukh3LjKiVhE9GwPTetB7spVe02Ddg6bCXJVUT4mjL84CS57mmV1iOtcQPp51ReVTCSZEgSTUm8-Uwp3q-T5VVwjdK7jTR6q3YdJx6jhzZ9FObvaEi87rezHMVvnPp51GEwahGUSRckucYvw-NoUBHf8Q1ZDrqSeEU9aGcsidfZW_QWvW1jTuIbOsOPcSkKS9m4_D40bLUKGQjUgmsr_w8xjLLeXN5xn2PKHAjnhGeE5xRV6r9-psNRPk-ifGD1PtMHdkmUBj-g2I75miyvMBdaL7wyOjAQ1ooDidcqXaB0NkmP4CxUA9ktrwhbUZLcFAXr2Au7W-IW1KqTI-wuG5_knRPo-DqqelBQlaOagc_R-G5zZ3_niwkXVOhI-mn0UJfP_XbWq6dR_ct46AoLxnbqt-AzrCZjuUGhr0g61BeH1Ug4BzbsAGtR4jEr6Gpcm3B-EM1cuafjgyEveEp4ih680y-iUuWJOi8gvbHBn0u6FaqCcj5asd6D_BGq0l45PBRPz-GAZHvvm5Co_Jbw253y-7aYS1MTfltVL8OfWWPNvwcDb5u2qgi_jaIkjkKIigMlC1YILeR-5hoh4bQeB4kNFkarqdubtio14UtPC6C18HIPJfL4I7isUj_hEUvWv2jlr1XtnoT88dg8ddvRv05pCSEWWFCHioUFC_1ptp28aXU0DRUWOpm4rxLOd4RAnWmthC5uSu7pq3AYr7aGkm6N7Zg9Pc_psehLdCVVmv6264pK1ED4LY-XRVYCl6zMARImU7FcLCBPJFuyVGayBLFkEBF-W1fKIiNVEH57rUTVcboP5hF--x2P-9bY2p38mMumITy-j6Ispg6gdlQbj5Ak0eNDTKxXsq2EpfAm6qaCU-fOIvRon_igS1qLAwZNSAk9xO2B_kUJv6IR9VZ0EI34gQ4VuuyEBqKT5TldYxrQg2lpqO22pqqPXDhPHnQJJYlvh-wVejcI7PyPAJMyC761Gnk-BcC66Yph5qx8FvoHnVHd1t8RqkIlukKDXlVVUSla1xk6ZKbSX3rFm3CeBqMwRfoiH6rMJN9M07GHF9C-FVUVHBZwDErqzVeIAG-NQItDLR5AITzbuL1ooC-rrAc6NoJZh3D8-DUOGNjhHDWtb1o_sOiRMJvA4QjbX4JggPJJqewWT0FoxJxy2pkcey6lXVfjJ2au_oEpz1FzN-JC9J7-HRLznL7Hs-CUAcPi0Vh0DMLHV-1H32gsrz_iAOFx96Z_mldMCQuNNWUrgbBV6EJf4nkyH9rD6e6u7Ws8bUIxm-OhpiTBbmY2a4Rzs0Y1UCkNJL4mnBetqrzS89qUbQUIhq2Wc_wgPPvFPpXn4c1HOSUU7W78hSrMSuVEUcHM7y2IEs_XR7O_WyEhAO2HpW8eq2vZ1k1oxU3rqTvUhamoFjU2ZjwD7VoL4ZzvxQvQUBA7IvUTLB66g2ktfVp9_5MaSx0M1eZFWaNr0J6-iNCf39__68Pm298PV4_3d_9287zBLXicvaGNUToUN-XHhoPRShUYlm52ADt3hrI3xhhbbkEuCiaypMiCRrgjXrmD6748WaV9MK4znmdHIiteN8Z5C6IOaLzGA9Jh8xVdoGOi84LTOIHPBD-3-pvaaVH9KXRZgXUj7l9RHmeErfgvWCWLlLBVTJFUzp2Zp1OKZAFpwlHRBT3LrExYEdOAQUHB1XarNDyIpvu5A7822sOb7xWVRjsf1I0JWyVnBUSMF3LxmYCV93YUQnj2mQ6DZ9CY9Kwt222-zCeirlpVlXAUchSNbV2QPw18ePoMWxKvP_XHzaBOFC8JWy3P2b7Ni2W0LSYKdWWz-_59qJ2PR1ff6RLelN49iGaaGItlQtgqOydPFjwplpgZ-VnVZBJtU7QiOneKtnleikX0hRVdl_KFCcGhEzuiIPDc6dkWkSxYiaTnzsJ2m0EsA2l8nrRI02ki9g3nqmkqJYU3_ePQr650-QyvVnk4SczHBlv-0K2vsDB8YNZvsn3hOOYWVnbcuLEhvwpjKnrCumfQna-wPe9HnC-ZvBhV_reZHBfuzU5JUT13o-0v8e3CmoXjGS3ORmDJcwgpkJwnZctienJE01SHXg33hwUoVXU40fEZdiEw6Wlgbq35CUNcegbfwE9MmVIHzoeeet11CnxNMVwh5H0eZxGOqtG5ilTKiEtWIOnyPGmUJxJJs3N1Lo4gkVPnlOCxd-1T9LvBYcq5VSmaMattq99F1LkPSfx1eq-0qA5OuQehxQ7Te_DJmrY6XO2VU5RMYyxYUX7ekm1U5L9tydPQSp2eTSQdFUz_x007Pe7O3WnnbYv9C47kH-rBlKCT3C30uj_hvOrv9NYMqTimVxwuQvjZ9ibmUc6Wv-29Rz06465uKsKzzt6hCYkilH62x4kjEJn8_5yFUZZwNOVsRxVzviiLd9EbhHyq-KmOvbgkDp47h0RFHCWswJGCnyuZRRwtyhjxjZ8rmUUcpTLnEyNcUym_0uWTNRKcu2q3W_RZ5nzZEbRa_UcLmwbbsAkOPEBt7GEgX9ORvoStaCu_KaECD1_vwRH9vww6_xfave_4P4Lmx5nguPbN45SJfeNXj3uYjEJ7f7aNjdk2TqeBfKiUfWz8g1D6zKDyv-e8o3b9jRXCr_P28KHyTrTvAHXaRXQlh2PPys9BZMFYBlFJcTp-bPxzq5FnPx9EOI3w7KuBSJZljk0KzycU4-sdKSwY3WyQbuO8sH5TD3IijiGMGR3H_fdMuhdnUZrDu3H6oryMyzzOxQVcRsuE5TFb5ouL_WUuUgYFh2KxjESeJkmcZLngIo4KmUmWXahLznjCUp6yPM6TxVzmiyLfZmmWgOC5kGTBoBaqmmPY5sbuLpRzLVxGiyRLo4tKFFC58L83zrvrRk6S6wt7GYb1ot05smCVct4dOXjlq_D_uof7u3CRnFz1jX9yjcP6L13mpizcD8rxanpyRQjaKaO7K9jw8KK11eVv37gGSx3ht72xL5f8PwMAAP__iNvEog">