[mlir][tensor] Implement TilingInterface for unpack op
The main issue of tiling unpack op is about incomplete tile. Since all the dimensions are orthogonal, discussing 1-d unpack case is enough. The core idea is to make the input slice have complete tiles. In this case, a larger unpacked tile will be created. We'll need an extract_slice op to shift and truncate the output. Take Nn_to_N as an example. Say that N=32, n=8, and tiling_size=15. The coordinates of second tile (i.e., result[15..31]) are [(1, 7), (2, 0,), (2, 1) ... (3, 6), (3, 7)]. The first row and the last row are incomplete in terms of inputs. It's impossible to represent an unpack op using the coordinates. Because the input has higher rank and the math computation of coordinate is using mod and ceilDiv. That's very tricky. To represent the unpack op, we have to complete the rows. I.e., the input coordinates would start with (1, 0); end with (3, 7). In this context, the tiled unpack produces a (3 * n) elements because there are 3 rows in total. Follow by a tensor.extract_slice op, we can get the actual result. If the tiling sizes are multiple of inner tile sizes, it is a perfect tiling case. In this context, the larger input and output is not needed. Reviewed By: chelini Differential Revision: https://reviews.llvm.org/D139362
Loading
Please sign in to comment