Lowering convolutions to matmul


I’m trying to optimize convolutions (any dialect) by converting it into a matrix multiplication, to do so it is required to implement a image2col operation within MLIR. AFAIK there is not something similar to im2col in the core dialects, right?

If I were to implement this, which core dialect should I use? The actual implementation of the algorithm is pretty similar to a convolution (think of 4 nested loops for 2d). Would this operation make the cut to be in a core dialect, like a linalg.im2col, tosa.im2col or something?

Is there any MLIR-based project working with convolutions that may have im2col or alternative forms for lowering convolutions?

@asaadaldien has something with im2col IIRC.
You should be able to achieve something relatively close with subview (but these will be small matmuls) or with packing (but it is at the level of linalg on tensors which is not fully hooked up end-to-end yet).

I think there is room for a notional transformation conv → matmul in Linalg by packing with redundancy.
It’s not available yet but contributions are most welcome! :slight_smile:

There is a pass that converts conv to matmul with im2col in IREE.

it packs the lhs of the matmul (the col-buffer) to stack memory. so this operation needs to happen after tiling the spatial dim of the conv to cut memory requirement. See this lit test

@asaadaldien @nicolasvasilache Thanks for the responses!

I’ll look into it.