[all-commits] [llvm/llvm-project] 83396d: [mlir][tensor] Implement TilingInterface for unpac...

[Han-Chung Wang via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 83396d85495914e17746d0ea6daa56eb35de3b39
      https://github.com/llvm/llvm-project/commit/83396d85495914e17746d0ea6daa56eb35de3b39
  Author: Hanhan Wang <hanchung at google.com>
  Date:   2022-12-16 (Fri, 16 Dec 2022)

  Changed paths:
    M mlir/include/mlir/Dialect/Affine/Utils.h
    M mlir/lib/Dialect/Tensor/IR/CMakeLists.txt
    M mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
    M mlir/test/Dialect/Tensor/tiling.mlir
    M utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

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