commit 01b9521f09062c857ad39592cb2c311c63dfcb11
author MaheshRavishankar <1663364+MaheshRavishankar@users.noreply.github.com> Wed Oct 09 10:32:34 2024 -0700
committer GitHub <noreply@github.com> Wed Oct 09 17:32:34 2024 +0000
tree 59e8d87c3ce31173f632097dec39241ca8bf13e6
parent 2963600a6f5a3d7a4a998ce8c7d4f9e46391c0cc

Add `TilingInterface` implementation of `iree_linalg_ext.custom_op`. (#18676)

This patch adds support for doing basic tiling (not fusion yet) of
`iree_linalg_ext.custom_op` by implementing the relevant methods of
`TilingInterface`.

It copies over a few methods from
`mlir/Dialect/Linalg/Utils/Utils.[h|cpp]`, modified to be able to handle
indexing maps with symbols. It uses a trick of converting these maps
(temporarily) to maps with just dimensions, for the rest of the logic to
compute tiles to be reused from Linalg operations.

In the end these copied methods could be merged back into main
(eventually when the op itself is upstreamed).

This patch also adds a `iree_linalg_ext.index` operation that mirrors
`linalg.index` operation, except it checks that the containing op is a
`iree_linalg_ext.custom_op` (instead of `linalg.index` op). The tiling
implementation handles the transformation of these
`iree_linalg_ext.index` operations similar to how the Linalg tiling
algorithm handled `linalg.index` operations.

---------

Signed-off-by: MaheshRavishankar <mahesh.ravishankar@gmail.com>