<table border="1" cellspacing="0" cellpadding="8">
<tr>
<th>Issue</th>
<td>
<a href=https://github.com/llvm/llvm-project/issues/56396>56396</a>
</td>
</tr>
<tr>
<th>Summary</th>
<td>
[MLIR] wrong vector shape after linalg vectorization pass leads to redundant vector.transpose
</td>
</tr>
<tr>
<th>Labels</th>
<td>
new issue
</td>
</tr>
<tr>
<th>Assignees</th>
<td>
</td>
</tr>
<tr>
<th>Reporter</th>
<td>
LeeOHzzZ
</td>
</tr>
</table>
<pre>
Hi,
I am having a strange issue while using MLIR.
I am lowering the following program with this command: `mlir-opt linalg_layernorm.mlir --linalg-fuse-elementwise-ops -test-linalg-codegen-strategy="anchor-op=linalg.generic register-tile-sizes=1,1,1,4 vectorize"`
```mlir
// This is a layernorm operation from the pytorch through torch-mlir
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0)>
#map2 = affine_map<(d0) -> (d0)>
#map3 = affine_map<(d0, d1, d2, d3) -> (d1, d2, d3)>
module attributes {torch.debug_module_name = "LayerNorm"} {
func.func @forward(%arg0: tensor<1x8x16x16xf32>) -> tensor<1x8x16x16xf32> {
%cst = arith.constant 0.000000e+00 : f32
%c16_i64 = arith.constant 16 : i64
%c8_i64 = arith.constant 8 : i64
%cst_0 = arith.constant 1.000000e-05 : f64
%cst_1 = arith.constant dense<1.000000e+00> : tensor<8x16x16xf32>
%cst_2 = arith.constant dense<0.000000e+00> : tensor<8x16x16xf32>
%0 = arith.cmpi eq, %c8_i64, %c8_i64 : i64
cf.assert %0, "mismatching contracting dimension"
cf.assert %0, "mismatching contracting dimension"
cf.assert %0, "mismatching contracting dimension"
%1 = arith.cmpi eq, %c16_i64, %c16_i64 : i64
cf.assert %1, "mismatching contracting dimension"
cf.assert %1, "mismatching contracting dimension"
cf.assert %1, "mismatching contracting dimension"
cf.assert %1, "mismatching contracting dimension"
cf.assert %1, "mismatching contracting dimension"
cf.assert %1, "mismatching contracting dimension"
%2 = arith.muli %c8_i64, %c16_i64 : i64
%3 = arith.muli %2, %c16_i64 : i64
%4 = arith.sitofp %3 : i64 to f32
%5 = linalg.init_tensor [1] : tensor<1xf32>
%6 = linalg.fill ins(%cst : f32) outs(%5 : tensor<1xf32>) -> tensor<1xf32>
%7 = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "reduction", "reduction", "reduction"]} ins(%arg0 : tensor<1x8x16x16xf32>) outs(%6 : tensor<1xf32>) {
^bb0(%arg1: f32, %arg2: f32):
%15 = arith.addf %arg2, %arg1 : f32
linalg.yield %15 : f32
} -> tensor<1xf32>
%8 = linalg.generic {indexing_maps = [#map2, #map2], iterator_types = ["parallel"]} ins(%7 : tensor<1xf32>) outs(%5 : tensor<1xf32>) {
^bb0(%arg1: f32, %arg2: f32):
%15 = arith.divf %arg1, %4 : f32
linalg.yield %15 : f32
} -> tensor<1xf32>
%9 = linalg.fill ins(%cst : f32) outs(%5 : tensor<1xf32>) -> tensor<1xf32>
%10 = linalg.generic {indexing_maps = [#map0, #map1, #map1], iterator_types = ["parallel", "reduction", "reduction", "reduction"]} ins(%arg0, %8 : tensor<1x8x16x16xf32>, tensor<1xf32>) outs(%9 : tensor<1xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
%15 = arith.subf %arg1, %arg2 : f32
%16 = arith.mulf %15, %15 : f32
%17 = arith.addf %arg3, %16 : f32
linalg.yield %17 : f32
} -> tensor<1xf32>
%11 = linalg.generic {indexing_maps = [#map2, #map2], iterator_types = ["parallel"]} ins(%10 : tensor<1xf32>) outs(%5 : tensor<1xf32>) {
^bb0(%arg1: f32, %arg2: f32):
%15 = arith.divf %arg1, %4 : f32
linalg.yield %15 : f32
} -> tensor<1xf32>
%12 = linalg.generic {indexing_maps = [#map2, #map2], iterator_types = ["parallel"]} ins(%11 : tensor<1xf32>) outs(%5 : tensor<1xf32>) {
^bb0(%arg1: f32, %arg2: f32):
%15 = arith.truncf %cst_0 : f64 to f32
%16 = arith.addf %arg1, %15 : f32
%17 = math.rsqrt %16 : f32
linalg.yield %17 : f32
} -> tensor<1xf32>
%13 = linalg.init_tensor [1, 8, 16, 16] : tensor<1x8x16x16xf32>
%14 = linalg.generic {indexing_maps = [#map0, #map1, #map1, #map3, #map3, #map0], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0, %8, %12, %cst_1, %cst_2 : tensor<1x8x16x16xf32>, tensor<1xf32>, tensor<1xf32>, tensor<8x16x16xf32>, tensor<8x16x16xf32>) outs(%13 : tensor<1x8x16x16xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32, %arg4: f32, %arg5: f32, %arg6: f32):
%15 = arith.subf %arg1, %arg2 : f32
%16 = arith.mulf %15, %arg3 : f32
%17 = arith.mulf %16, %arg4 : f32
%18 = arith.addf %17, %arg5 : f32
linalg.yield %18 : f32
} -> tensor<1x8x16x16xf32>
return %14 : tensor<1x8x16x16xf32>
}
}
```
and got the following:
```mlir
module attributes {torch.debug_module_name = "LayerNorm"} {
func.func @forward(%arg0: tensor<1x8x16x16xf32>) -> tensor<1x8x16x16xf32> {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant dense<2.048000e+03> : vector<1xf32>
%cst_1 = arith.constant dense<0.000000e+00> : vector<1xf32>
%cst_2 = arith.constant dense<2.048000e+03> : vector<1x1x4x1xf32>
%cst_3 = arith.constant dense<0.000000e+00> : vector<1x1x1x4xf32>
%c1 = arith.constant 1 : index
%c4 = arith.constant 4 : index
%c8 = arith.constant 8 : index
%c16 = arith.constant 16 : index
%true = arith.constant true
%cst_4 = arith.constant 1.000000e-05 : f64
cf.assert %true, "mismatching contracting dimension"
cf.assert %true, "mismatching contracting dimension"
cf.assert %true, "mismatching contracting dimension"
cf.assert %true, "mismatching contracting dimension"
cf.assert %true, "mismatching contracting dimension"
cf.assert %true, "mismatching contracting dimension"
cf.assert %true, "mismatching contracting dimension"
cf.assert %true, "mismatching contracting dimension"
cf.assert %true, "mismatching contracting dimension"
%0 = linalg.init_tensor [1] : tensor<1xf32>
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<1xf32>) -> tensor<1xf32>
%2 = vector.transfer_read %1[%c0], %cst {in_bounds = [true]} : tensor<1xf32>, vector<1xf32>
%3 = scf.for %arg1 = %c0 to %c8 step %c1 iter_args(%arg2 = %2) -> (vector<1xf32>) {
%20 = scf.for %arg3 = %c0 to %c16 step %c1 iter_args(%arg4 = %arg2) -> (vector<1xf32>) {
%21 = scf.for %arg5 = %c0 to %c16 step %c4 iter_args(%arg6 = %arg4) -> (vector<1xf32>) {
%22 = vector.transfer_read %arg0[%c0, %arg1, %arg3, %arg5], %cst {in_bounds = [true, true, true, true]} : tensor<1x8x16x16xf32>, vector<1x1x1x4xf32>
%23 = vector.multi_reduction <add>, %22 [1, 2, 3] : vector<1x1x1x4xf32> to vector<1xf32>
%24 = arith.addf %23, %arg6 : vector<1xf32>
scf.yield %24 : vector<1xf32>
}
scf.yield %21 : vector<1xf32>
}
scf.yield %20 : vector<1xf32>
}
%4 = arith.divf %3, %cst_0 : vector<1xf32>
%5 = linalg.fill ins(%cst : f32) outs(%0 : tensor<1xf32>) -> tensor<1xf32>
%6 = vector.broadcast %4 : vector<1xf32> to vector<1x1x4x1xf32>
%7 = vector.transpose %6, [3, 0, 1, 2] : vector<1x1x4x1xf32> to vector<1x1x1x4xf32>
%8 = vector.transfer_read %5[%c0], %cst {in_bounds = [true]} : tensor<1xf32>, vector<1xf32>
%9 = vector.extract %cst_1[0] : vector<1xf32>
%10 = scf.for %arg1 = %c0 to %c8 step %c1 iter_args(%arg2 = %8) -> (vector<1xf32>) {
%20 = scf.for %arg3 = %c0 to %c16 step %c1 iter_args(%arg4 = %arg2) -> (vector<1xf32>) {
%21 = scf.for %arg5 = %c0 to %c16 step %c4 iter_args(%arg6 = %arg4) -> (vector<1xf32>) {
%22 = vector.transfer_read %arg0[%c0, %arg1, %arg3, %arg5], %cst {in_bounds = [true, true, true, true]} : tensor<1x8x16x16xf32>, vector<1x1x1x4xf32>
%23 = arith.subf %22, %7 : vector<1x1x1x4xf32>
%24 = vector.extract %23[0, 0, 0] : vector<1x1x1x4xf32>
%25 = arith.mulf %24, %24 : vector<4xf32>
%26 = vector.reduction <add>, %25, %9 : vector<4xf32> into f32
%27 = vector.insert %26, %cst_1 [0] : f32 into vector<1xf32>
%28 = arith.addf %27, %arg6 : vector<1xf32>
scf.yield %28 : vector<1xf32>
}
scf.yield %21 : vector<1xf32>
}
scf.yield %20 : vector<1xf32>
}
%11 = linalg.init_tensor [1, 8, 16, 16] : tensor<1x8x16x16xf32>
%12 = vector.broadcast %10 : vector<1xf32> to vector<1x1x4x1xf32>
%13 = arith.divf %12, %cst_2 : vector<1x1x4x1xf32>
%14 = vector.transpose %13, [3, 0, 1, 2] : vector<1x1x4x1xf32> to vector<1x1x1x4xf32>
%15 = arith.truncf %cst_4 : f64 to f32
%16 = vector.broadcast %15 : f32 to vector<1x1x1x4xf32>
%17 = arith.addf %14, %16 : vector<1x1x1x4xf32>
%18 = math.rsqrt %17 : vector<1x1x1x4xf32>
%19 = scf.for %arg1 = %c0 to %c8 step %c1 iter_args(%arg2 = %11) -> (tensor<1x8x16x16xf32>) {
%20 = scf.for %arg3 = %c0 to %c16 step %c1 iter_args(%arg4 = %arg2) -> (tensor<1x8x16x16xf32>) {
%21 = scf.for %arg5 = %c0 to %c16 step %c4 iter_args(%arg6 = %arg4) -> (tensor<1x8x16x16xf32>) {
%22 = vector.transfer_read %arg0[%c0, %arg1, %arg3, %arg5], %cst {in_bounds = [true, true, true, true]} : tensor<1x8x16x16xf32>, vector<1x1x1x4xf32>
%23 = arith.subf %22, %7 : vector<1x1x1x4xf32>
%24 = arith.mulf %23, %18 : vector<1x1x1x4xf32>
%25 = arith.addf %24, %cst_3 : vector<1x1x1x4xf32>
%26 = vector.transfer_write %25, %arg6[%c0, %arg1, %arg3, %arg5] {in_bounds = [true, true, true, true]} : vector<1x1x1x4xf32>, tensor<1x8x16x16xf32>
scf.yield %26 : tensor<1x8x16x16xf32>
}
scf.yield %21 : tensor<1x8x16x16xf32>
}
scf.yield %20 : tensor<1x8x16x16xf32>
}
return %19 : tensor<1x8x16x16xf32>
}
}
```
The issue is that during the linalg vectorization pass, it creates a `%cst_2 = arith.constant dense<2.048000e+03> : vector<1x1x4x1xf32>`, which is different from the tile size I specify (`1,1,1,4`), and this leads to some redundant transpose operations around here:
```mlir
...
%12 = vector.broadcast %10 : vector<1xf32> to vector<1x1x4x1xf32>
%13 = arith.divf %12, %cst_2 : vector<1x1x4x1xf32>
%14 = vector.transpose %13, [3, 0, 1, 2] : vector<1x1x4x1xf32> to vector<1x1x1x4xf32>
...
```
Is this pass supposed to create a constant vector that is different from the specified tile size? Also, is there a way to optimize away these redundant transposes?
Thank you!
</pre>
<img width="1px" height="1px" alt="" src="http://email.email.llvm.org/o/eJztW1tzo7gS_jXkRWUX4mb84IdJslM7VXPOqTq1T-fFJYOwtYuBRWKczK8_LQkwGIFxksns1mbK43BRt1p9-5qW2eXx8-ZXZjkPlv1o2Z_09xdEjuhAvrFsjwjioiTZniLGeUXR6cBSiiou7_3r65f_LjskaX6ipbwhDhQleQrn8qwo830Jt09MHOAW4yjKj0eSxZb7CVmBfUxZucgLgVKWkXS_TckzLbO8PC7lHbRY6OuLpOJ0QVN6pJk4MTjOC44WgnLRjIjymO5ptpAiC7p_ttxHy3FIFh1yOQOc6oFLGASSRqike8YFLRcCVrXg7DvlMAiDPpr_HvpGI5GXcAtYgbS1ouBAfaSI9SXnM3zQb3KB8CGoXQfKCwoCsTxDSZkflXqKZ2AaSX2UebWHv_Js0eXmHklhIxAHkSRhGd3CueU-WE4Y2yAYirH6dtS3azlrtLDcX9D4fbjb5Y1fzPuSkzPKaZrOfZkEo-s65nEFzkmEKNmuAsdA1upeKXYZ01213-oB24wcqZoaLPpVGunfYCRp3dWjpNDMEEqqLFrKL2R5dpKXJ1LGML_l-KTc29J5Bc14XoLY-Cl8woH8JK4j5WnkHR3RnQiBIH7EhVZHCWGyjPKMC5IJZC9t9Q-8796W7vAJSQY9ShxsWeCZqHGgKOBunyIcIwjN47nY2kb-jXQL29eyGSixiTIGxVCpl976lGK6ir1Q6yVrZ4q1_WLWvbUeC4bon9LbWs31ToYai5Il4ZyWQvHSg50j40ciooPMiCAqZKhIyOOYQT7jkBukA_7FWAAhHleFdrv-2RVl4Nev5IPFa1kAYTdujlXKhq49Zk645xqInRl03YTDmciTouGmxgIEDlKbr0hq0GYZE1sdvsjy77HlP6KLLGyI5aDLImFpiljGdRrXKVcnVEjYeSXqG_4IX0NWN8y46s7YVBqQ7lkW0ycwi4Q6rvHHv69xXqtPwTIsS54xIWuGvNyK54J2hjsFKUma0lTaVJu_pHEVidrM8y7BHAB2rSYkoF2ueYhoZwUF4wrq45r_y25nt3PgVtvKWeCKc9Y_HJ0Jde7xOx5D4jhpiVp6PERE1Kj-mdE0bvlc4iYsf4Ytw9tt6Zxt6dxiywuTrMZVPMtRf5QdYvatsUOdfnzvBxth_e4hjO1XxvDPD-faNuG1sH645mPrt_ex5oo72-t4tbv0OsnR5HiSNOhDVKL51YQmT9RkK3O-cRvCYJajr17s6Bj_tHSDB_n_n5tvsPPzzID_omYQJTwIJ92HQfWwNyjahvHXCSQ8MwKhmD0sS_5nXeX-8LBzJ8tMkDmUXziov4d158SzJPbeDkraQ9d4aN-ONbdemQCaxrbto4B89O8cOy9ComsXJ5hMFbDYvV7uvh26NVe8wRV_cCX4KZgoxZ2Dii1h0FnVGGFoyAJ41Vn8rLgO58X1WAiWVFRl1kbijLAF7nWXsj1our36lGQx2uei3-ZujWVoDf-9G5PGzp9GAJXD3qqNOdZkbPp5ztL2woaH2_TzdHt-LLVf7T-am4TXmU52Hq9Kip-8p1HW7qvk1cxNrI16wGN2NLaHvbHRoWl0ODa6l4sGzerBeKg9qIlAXh9oz9wHn-xT93pqiumrO3MfXD64XG_D9roML2tv9p4YZ_ZGJh71ZlTLOvPpjLOUe8M8oeW2pETDtao0ATTqWrSRQ9a7211eZXFbkCrN6YJyRJ6HK4lYZ0oO9kik0tqO4GONW_BwopMTF7Soc6Asjrcw7lzBOg2F091mHM58WRJqbdgGGdyhDJDcJoXwGhLd37xNDi0JNkjiT0viGSQJOpJ4t0uiZZl2ElWmNH5y7uV2S-dOhTzPkWTlb_prdLDhU8NVBD2vze2uDQpiwbZtMw5uPUCpWzOtNVE_RKoi323CemRCaaUJpz9L0cW5prh2OnoLrhUy-p90l7bWdrw5RG1ZbOaBr_O44NCnt69WYF1qq7-v1HSL3LPLbK8zvNhoepc0GnS9aFfmJI4IF-f21kDeC88YLSFXg9Arck7VlEop_r1Sjgq82i9NLnnmPph4pMAMJ2Pef09gWHdFoU8Kic89Cf_eNix5fBvgzQAm_ACYD4CZCzD9Jk_Tjav3BW_g5o3EAoCFDIQmFZhC4hrrbjeq6Qw5zU8ILsFkilEvGU6AadOuWptZQ9Ie9qObOXp5EbJ7Xb87QQcqMOomB2CkOc4CZFO3y1m9EpDDvyMg97ex3rqf3ov8HnDiEUnnIifuBl5TSfT62c7sRg72BgmqgWHs_lgcnti38Ub3bay2Q2zUbNuqnS2EaS8VN7kBD4JhnFFo2hKamQXl0PXbYjjGXciav4HwPmB-izzvA-q3SfQB7i8B9wsEbn-vMECPWwC9xTCvkwDdWzkGRlOeYAbaQ3S183WDUV9jyjHx-5uNYxhkgsjB7-GmiGfB9SxeM2B7Dqh2uXS2ywY__Hnhdpn-_u3QvELCOBIHIlBctS-L6FKhfdVCvytREM71djaKIPblthmR74t0sPiNdmAC5XKnA4sOUrqYJeCkFNi1r2vI90OQfD8EfUG8oBFLnmVyA8reyyKK01oyk3uE6l2XFLIWl1mU50eKZG2bxXrzoikJ2rdDYH2l9Gd0gNmn9hOXy-VHTTSvJmpVZXTKL1wbSboa4lUhxYklT-1w4G-tZ-lZtOeafUT7BZMMGm-x3M_oU8pz5cZyLiABpifyLCfJC8GO0qeIunCg3OggHLj0I4lkf6DnvLIcfBdv3HjtrsmdYCKlG9CgfCtLau1U5lkTUogfSAHzJIDb48F29tWzFJfGuqvKdHMQouDSQdUrT3twjWoHMXiEkzT91vxZFGX-O5DDqQp8CObPfuCug7vDJox2ztqlCSbY9Z3dLnQxpokdhCSIvJW3ukvJjqZ8o3-3ktGTzh36Jyh3bOPYjmOv7MC23cCxl8AAr5N16Me7cBVh3_JseiQsXUo5lnm5vys3SqRdBTWLZ6eMC36-CWtn-4wq9Un-pBKHvNx8pfQ_v37__r87NfdGyf5_mf86ww">