MLIR: add new unit dims folding to improve vectorization.

[liamsemeria]

Motivation

In order to vectorize conv2ds properly, xtc relied on a combination of unit dims folding and the fastmath<fast> flags on the arith ops in the conv2d linalg.generic. Both of these are required for conv2d to get fma instructions later in the pipeline, during the mlir llvm dialect stage.

Comparatively for matmuls, instead of being fastmath ariths that get raised to an llvm fma op and the end, it roughly lowers like this:
linalg -> vector.contract -> vector.outerproduct -> vector.fma -> llvm.fma

This new version of folding makes the lowering the same for both matmul and conv2d, does not rely on the fastmath attribute, and allows for better vectorization on matmul linalg.generics and the memref dialect.
There is more info on this new transform op in the pr for it op in the xtc-mlir gitlab.

Description

The new folding transform op is used on both memref and tensor dialects if xtc_transform is installed, if it isn't installed then it it resorts to folding all unit dims for tensor and memref. Resorting to folding for memref instead of the previous no folding is because in most cases for memref in this llvm version, folding is still more beneficial than not.

I kept the fastmath attribute in the mlir conv2d op, so that vectorization would still work if xtc transform isn't installed.
While using this new transform op, performance is the same wether or not you keep the fastmath attribute for both tensor and memref dialects.

Discussion

• Is this optional behavior a good idea? yes
• If it is there a different way to implement it that is preferable? no it's fine

[codecov]

Codecov Report

❌ Patch coverage is 62.50000% with 6 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/xtc/backends/mlir/MlirBindingsExtensions.py	53.84%	4 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

[guillon]

Fine for me.

For the fallback it's ok I guess, the initial idea was to keep additional xtc mlir modules optional for peaople having their own local MLIR build and that do not need extar stuffs.
It seems fine w.r.t. this use case.

Just for the warnings, i'm afraid it may raise too much, see the question.

[guillon]

I don't know the warnings module, can you explain how to disable it and/or how it is preferable than a classical loggimg.logger

[liamsemeria]

I initially assumed that using the logger for warnings would require more set up in other files but I was incorrect. I changed it to use logger instead.

[liamsemeria]

warnings module outputs less warnings than logger, but both are outputing at least once per loop-explore iteration so I'll find a way for the logger warning to output once.

[guillon]

Actually I'm not necessarily against a warnings module which possibly avoid repeatedly output the same warning again and again. I let you take the best and most common approach

Just, I do not know how to configure it, while for loggers, I just have to set error level at the root logger and it works for all modules at one.

If warnings is a common pythonic way now, I'm ok with this, but how to disable it?

[liamsemeria]

I think the warnings module was still emitting once per iteration so it still had the same problem with loop explore. I resolved it for the logging module with a lock and attached it to the MlirBindingsExtensions.

The warning module gets disabled by environment variables or python cmd line flags so it seems easy to disable, if this warn once thing is too convoluted we can switch back to the warnings module.

[guillon]

Ok, it's fine, thks a lot.