search

iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.
http://iree.dev/
Apache License 2.0
2.47k stars 551 forks source link

[Im2col] Add decomposition pass for iree_linalg_ext.im2col #17728

Closed Max191 closed 4 days ago

Max191 commented 6 days ago

This PR adds a pass for decomposing the iree_linalg_ext.im2col op into serial loops of extract_slice->copy->insert_slice. The pass tries to keep the innermost dimension of the iteration space un-tiled when the inner slice is contiguous, but otherwise falls back to serial loops of scalar slices. The outer dimensions (B and M) of the iteration space, however, are always tiled to 1.

Max191 commented 5 days ago

based on https://github.com/iree-org/iree/pull/17671. Just review the last commit

EDIT: rebased now

Max191 commented 5 days ago

Do we implement the AggregatedOpInterface method (i.e., decomposeOperation) for the op? So like we can move the decomposition logics to AggregatedOpInterfaceImpl.cpp.

Oh, I didn't know this existed, thanks for the suggestion!