Performance Gap: conv+pool is not fused into one linear operation

[Zohaib58]

HEIR's pipeline does not fuse linalg.conv + linalg.pooling into a single linear operation (BSGS matrix-vector product). The LinalgCanonicalizations pass converts linalg.pooling_nchw_sum into a separate linalg.conv_2d_nchw_fchw, so each pool consumes its own multiplicative level. In contrast, a hand-written CROSS implementation fuses conv+pool into one matrix, saving ~2 levels for AlexNet-Tiny (depth 9 vs depth 7)

[asraa]

Thank you! This seems like a fairly straightforward rewrite pattern to write in linalg-canonicalizations (https://github.com/google/heir/tree/main/lib/Transforms/LinalgCanonicalizations) or linalg-fuse-linear-ops (https://github.com/google/heir/tree/main/lib/Transforms/LinalgFuseLinearOps) (perhaps we should generalize the latter to LinalgFuseOps)

Would you be up for that? Also happy to write!

[j2kun]

I will take this one on today

[j2kun]

@Zohaib58 I think the right approach depends on exactly what IR you see as input here. Is it specifically a linalg.pooling_nchw_sum op followed by a linalg.generic doing a division? Or a sequence of linalg.conv_2d_nchw_fchw followed by linalg.pooling_nchw_sum?

[j2kun]

ah, I think the problem is that before we have the opportunity to fuse, we rewrite pooling_nchw_sum as a conv2d, and then conv + pool becomes conv + conv.

I'm trying to decide if the best place to do this transformation is in canonicalizations or elsewhere...

Options:

1. Convert in linalg-canonicalizations directly to a matvec with the expanded filter. This would work but bloats the IR and duplicates work with the layout system, and it precludes the layout optimizer from picking a different kernel (right now there are no better options...)
2. Make a new op tensor_ext.fused_conv_pool that can be passed through to the layout system and treated as normal. Then fuse the ops into this new op.
3. Extend the layout system to handle multiple ops in sequence (messy)
4. Convert both the conv2d and the pool to linalg.generic, merge them into a single linalg.generic op, and hook off that in the layout pipeline (probably too much extra infra work for now).

[j2kun]

I think I have a handle on what to do now. The conv + pool can be fused into a single conv with a larger filter, which is what the CROSS code appears to do as it flattens.

I will implement this, but @Zohaib58 if you have a moment to post the specific (subset of the) IR you want to lower, I can include that as a test case.