search

iree-org / iree

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

[LinalgExt] Add iree_linalg_ext.im2col op and verifier #17644

Closed Max191 closed 2 months ago

Max191 commented 3 months ago

This PR adds a new iree_linalg_ext.im2col op representing the im2col transformation for convolutions. The PR contains the op definition and the verifier.

Max191 commented 2 months ago

Can we add few tests to cover the size(m_offset) > 1 case and size(k_offset) > 1 case?

I am actually just going to force these offsets to be size 1 for now. Size greater than 1 means that the matmul consuming the im2col op has multiple M or K dimensions. I am only implementing the decomposition and tiling for M and K size 1 cases right now, so I realized I should probably just restrict this. These cases may actually show up when we have data tiling for convolutions on GPU, but that will take a bit of time, so I am going to wait until later to support multiple M and K dims.