<table border="1" cellspacing="0" cellpadding="8">
<tr>
<th>Issue</th>
<td>
<a href=http://email.email.llvm.org/c/eJy1V11v4jgU_TXhxSoKToHywMNMmWpH2pFWo3kfmdgh3iZ21nag7K_fc50EaIG209FWNCTO9f0891yztnK_vG-dUyZUe5bM0qCMt26cCx_wxISRTAfPhPc21yIoyaSuIaOtYdoU1tUixHvPjA2scbYRmygXSmfbTckqbUS1YbbxY_bgrIEF6dlOh5LJvRG1zpkvRaOYqrWPekMpAnMKDuHh259fv7Ngj5o7ac8Sfhesy8ubutIu4fcHbXVZWXpWIU_4gvm2KJRDBExUFbnBhFMsOBW9xPKwT20V4ihgHjYCwsoHU7TB6xC02TCnN2VgRj0F8gqydZJ9StJVkg7XWdp9ol9xicHZKeczlmSrPh9jbXT42WWbJdPPE_J4OuH01V0znkxX2PGJdVJJdj95gsQTPhl_KvA--_JM_zzqP6ng0eyrSmIgw-ske3ii_6tW7k6j2CijHBKVzD9rI9UTMvSzFsgxySAsURTaKFoi1fxOprFSMVoZw5QZFemG3Lj-nnzA_Qe18ROpQdt0RY86KCeAop9hHyHVOZ1w3ggHtKgKtyT3gRUYmK_QIh5-xLxNn5ehzzF9HfMM320bDlvm51vOSoMtSP6hQtMv63Xa7RduM6H9JBj9oxV-WFnQHVYf8EHvosGUVLlCozs_qGNDnfdaVbJXMWGDjsEoAo0pfweIDt1x2jJ_KDSYDgmfe_Zo7K6nAGrEvkNK9Gnfk5USkhrxwEOe2YKoKwKbSGvdBrZWuWi9OmvmC0TVUKeQ3Dn90aptwEW1QCUHphif-v6jhEbRNEo4zyIpNrGC92xX6rxkYiDAXBhrdC4q_S_oKLctElpYXLShBmwbSeQWYyYwHlnlCp_stARfXeh4lCg9Bw4q8Vq7Zy9a_UV_j_F3BPNg-gI2syMuO1U371X6ltsHjRch9DVQgkHx3vbVoGKgmaga3bQRUmqaVaIaEHEywN5ZsOMUahu23rPabgmLVLWY_m4kQBgj1TfWRKRGSLxZTxoIrwwISADxad_JuI18lp3NiGsFxSZk4xJ3xzIcaaf3430qjXDO7i6hsDP2ipYzGPIXBm6em_r_8xM1dv6fx_N2Wt4M6BeLMLhyvSMupQlwBojJgnBA_TivG83UPy_ygtucR91xbA_7ccxTLhzU9GMNmzyrWzjTr_LD3EDIa3VKpqryalcSn6O1cHwU1d5rBPVAp7GPFJNfL-Nv5vdthReZ5gf6O2oFj4AHBU0U4m8MG0FpQvbP6YWJrdCVWFdRqLE46OJ-fGHk0DQxNHTiCFQoIchHm-CsbDGdnymlI46uDwfYOIM72_Ek4Whw2S1tCoNs71oU7fdB9tlEQ7FynOhu02-iVL78LrYao9M8glJHcpnJRbYQI9GG0rrlWpmtMPpx1LpqWYbQxLkVsbGB6-16nNsaD1W1Hb5ugJa_VY5p_4DzfqsIHNMsS_moXGapUsVcSTmdz6Z3-WK6uF2IWzXJZSEkn2cjpBAIWxJEpqvR75vUS55ynk4mCz7J5pPZOL3NinmeytviLp_MZwJpoOFfjUnP2LrNyC2jynW78XhZaR_88SUaSG-MUoOHQYcqPtDPGALxX4cB0qOd6u4v_VgaRVeX0c__AHFNJ1Y>53302</a>
</td>
</tr>
<tr>
<th>Summary</th>
<td>
[MLIR] Propagate tensor casts through linalg ops.
</td>
</tr>
<tr>
<th>Labels</th>
<td>
</td>
</tr>
<tr>
<th>Assignees</th>
<td>
</td>
</tr>
<tr>
<th>Reporter</th>
<td>
benvanik
</td>
</tr>
</table>
<pre>
Currently `tensor.cast` and its associated dimension information is not propagated through linalg ops. Frontends with dynamic shape emission that rely on MLIR to propagate shapes (torch-mlir, dynamic mhlo, etc) suffer as all ops are treated as dynamic even if the static shapes are sitting right next to them:
```mlir
%226 = linalg.init_tensor [1, 512, 12, 32] : tensor<1x512x12x32xf32>
%227 = tensor.cast %226 : tensor<1x512x12x32xf32> to tensor<?x?x12x32xf32>
%228 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d1, d3)>], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%225 : tensor<?x12x?x32xf32>) outs(%227 : tensor<?x?x12x32xf32>) {
^bb0(%arg1: f32, %arg2: f32): // no predecessors
linalg.yield %arg1 : f32
} -> tensor<?x?x12x32xf32>
```
Here it's known that the tensor has static leading dimensions of `1x512` but because the static shape is not propagated past the `tensor.cast` the op remains dynamic.
This appears on inputs, which a linalg canonicalizer could fold in to update the types:
```mlir
%widen = tensor.cast %arg0 : tensor<?x2x3xf32> to tensor<?x?x3xf32>
linalg.generic .... ins(%widen : tensor<?x?x3xf32>)
->
linalg.generic .... ins(%arg0 : tensor<?x2x3xf32>)
```
It can also appear on outputs with additional static information, which a linalg canonicalizer could propagate up by moving the cast to the corresponding input:
```mlir
%init = linalg.init_tensor [%dim0, %dim1, 3] : tensor<?x?x3xf32>
%out = linalg.generic .... outs(%init : tensor<?x?x3xf32>
%narrow = tensor.cast %out : tensor<?x?x3xf32> to tensor<?x2x3xf32>
->
%init = linalg.init_tensor [%dim0, %dim1, 3] : tensor<?x?x3xf32>
%init_cast = tensor.cast %init : tensor<?x?x3xf32> to tensor<?x2x3xf32>
%out = linalg.generic .... outs(%init_cast : tensor<?x2x3xf32>
->
%match = arith.cmpi eq, %dim1, %c2 : index
assert %match, "dims must match" // to be propagated elsewhere for analysis/etc
%init = linalg.init_tensor [%dim0, 2, 3] : tensor<?x2x3xf32>
%out = linalg.generic .... outs(%init : tensor<?x2x3xf32>
```
The general idea is to use as much static information available as possible. `tensor.cast` is an op that either introduces information (dim dynamic -> static) or removes it (dim static -> dynamic).
/cc @MaheshRavishankar
</pre>
<img width="1px" height="1px" alt="" src="http://email.email.llvm.org/o/eJy1V9tu4zYQ_Rr5hYghU77ED37YTTboAl2gWOx7QIm0xUYiVZKy4359z1CS7cR2kqZo4MgSNZzbmTlD51buV3etc8qEas-SeRqU8daNC-EDnpgwkungmfDeFloEJZnUNWS0NUybtXW1CPHeM2MDa5xtxCbKhdLZdlOyShtRbZht_Jg9OGtgQXq206Fkcm9ErQvmS9Eopmrto95QisCcgkN4-PH7958s2KPmTtqzhN8G64rypq60S_jdQVtdVpaeVSgSvmS-Xa-VQwRMVBW5wYRTLDgVvcTysE9tFeJYwzxsBIRVDKZog9chaLNhTm_KwIx6DuQVZOsk-5Kk90k6XOdp94l-xSUGZ2ecz1mS3ff5GGujw2OXbZbMvk7I49mE01d3zXgyu8eOL6yTSrK7yTMknvHJ-PMa77NvL_Qvov4TBI9m31QSAxleJ9nDM_1ftXJ7GsVGGeWQqGTxVRupnpGhx1ogxySDsMR6rY2iJVLNb2UakYrRyhimzAikG3Lj-nvyAfef1MZPpAZts3t61EE5gSp6DPtYUp3TCeeNcKgWVeGW5D6xAgOLe7SIhx8xb7OXMPQ5pq9jnuG7bcNhy-J8yxk02ILkHxCafcvztNsv3GZC-0kw-kcr_LCypDusPuCD3kWDKakKhUZ3flDHBpz3WlWyVzFhg47BKAKNKf9AER2647RlflNoMB0SvvDsydhdTwHUiH2HlOjTvicrJSQ14oGHPLNroq5Y2ERaeRtYrgrRenXWzBeIqqFOIblz-qNV24CLagEkB6YYn_r-q4RG0TRKOM8iKTYRwTu2K3VRMjEQYCGMNboQlf4bdFTYFgldW1y0oQZsG0nkFmOmYjyyyhU-2WkJvrrQ8YAoPS8cIPFWu2evWv1Vf4_xdyzmwfSF2syOddmpuvmo0vfcPmi8WELfAyUYFO9tjwaBgWYiNLppI6TUNKtENVTEyQD7IGDHKdQ2LN-z2m6pFgm1mP5uJEAYI9U31sRKjSXxLp40EN4YEJBAxad9J-M28ll2NiOuAYpNyMYl7o4wHGmn9-NjKo1wzu4uVWFn7A0tZ2XIXxm4eWnq_89P1Nj5fx7P-2l5N6B_CcLgyvWOuJQmlDOKmCwIh6ofF3WjmfrrVV5wW_CoO47tYT-OecqFg5p-rGGTZ3ULZ_pVfpgbCDlXp2SqKq92JfE5WgvHR1HtvUZQD3Qa-wyY_DqM_zG_7yu8yDS_0N9RK3gEPChoohB_Y9gIShOyf04vTGyFrkReRaHG4qCL-_GFkUPTxNDQiSNQAUKQjzbBWdliOr9QSkccXR8OsHEGd7bjScLR4LJb2hQG2d61KNrvg-yLiQawCpzopukPUSpf_hRbjdFpnkCpI7nK5DJbilHQoVIrQEXnc0LnjwMz9jBSQP7Sr4BR66pVGUITZ1ysow3CbPNxYWs8VNV2-LpBZf2pCpwMHvDboFVUSLMsS_moXM0WUt6mxWK6zqaT5VLKYipzNQXL82y-VGKEdKMayUf4N9IrnnKeTiZLPskWk_k4nWbrRZHK6fq2mCzmAgHTmK_GZHhs3WbkVtGHvN14vKw04jm-RKvojVFq0C_aUFq3ypXZCqOfRtHfVXT2HwscEg8">