[mlir] Fix `lower_unpack` when dynamic dimensions are involved (#68423)
When lowering `tensor.unpack`, we need to use the sizes of the destination tensor in the final `tensor.extract_slice` operation. Prior to this patch, when the destination tensor had dynamic dimensions, we would compute them from the result of the `tensor.unpack` operation instead of its destination argument. This would produce invalid IR because the `tensor.dim` operations would need to appear before the `tensor.extract_slice` operation, but the input of the `tensor.dim` operations would consume the final result of the lowering of `tensor.unpack`, which happens after the `tensor.extract_slice` operation. In other words, the definition wouldn't dominate its uses. I.e., we were generating: ``` %dynDim = tensor.dim %defLater, ... <-- %defLater defined below %res = tensor.extract_slice ..., %dynDim, ... %defLater = linalg.copy (ins %res) ``` Note: I checked the implementation of `lower_pack` and the code is correct as far as I can tell.
Loading
Please sign in to comment