[all-commits] [llvm/llvm-project] 7050ff: [mlir] Fix `lower_unpack` when dynamic dimensions ...

[qcolombet via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 7050ff4615d7341eab7fca821b061b3fe3dd4bd8
      https://github.com/llvm/llvm-project/commit/7050ff4615d7341eab7fca821b061b3fe3dd4bd8
  Author: qcolombet <quentin.colombet at gmail.com>
  Date:   2023-10-06 (Fri, 06 Oct 2023)

  Changed paths:
    M mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
    M mlir/test/Dialect/Linalg/transform-lower-pack.mlir

  Log Message:
  -----------
  [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.