Transforming linalg with multiple generic operations

[giuseros]

Hi all,
Following my previous post I now have a program with two generic operations in linalg:

// %0 = %beta * %C
%0 = linalg.generic ...
// %1 = %alpha*%A*%B + %0  = %alpha*%A*%B + %beta*%C
%1 = linalg.generic ...

The reason why I am implementing matmul through linalg.generic is because I want to have the ability to:
a) Changing the maps to implement trans(A)*B
b) Fusing the multiplication by %alpha directly inside the main kernel
c) Making the framework generic without using specific named linalg ops

Issue

When I try to vectorize the linalg program it looks like only the first linalg.generic gets vectorized, while the second is left only tiled. Is this supposed to happen? I also tried to tile first and then to have a separate vectorize transform, but as long as I have the element-wise operation in the middle, the vectorizer seems to not pick up the second operation.

General question

Those two operations are quite different and need different transformations (fusing everything for %0 + the normal matmul transformations for %1). Is it possible to "name" those operations and apply a transformation only on the specific named Op? For instance, we have a anchor-func, anchor-op in many transformations. Can we also have a anchor-name + a name attribute in the linalg.generic operation?

I am not sure, but I saw this issue: #149 and I was wondering if this was going in this direction. @apaszke , @Mogball could the strategy dialect be used in this case?

Thanks,
Giuseppe

[Mogball]

You could name the op directly with an attribute and have your transformations hook on to that.

In the rewrite system going forward, we're going to move away from relying on attributes and hooking on to them. Instead, you would want your transformation to return handles to the two generic ops and then directly apply additional (potentially different) transformations on those two handles.

[giuseros]

Hi @Mogball ,
Thanks for the answer. Just to understand, let's say I want to flatten the element-wise operation and tile the matmul, how do I do this without naming the operation? Is it through a pattern matching system?

[Mogball]

Transformations on operations do not have to be through the pattern rewrite system. You can write a pass and do whatever mutations you want, however you want. But the pattern rewrite system provides structure and reusability. If you can write a pattern to match only the element-wise op and another one that only matches the matmul, then you can use patterns.

[giuseros]

This is the feature that @nicolasvasilache added in the latest commits, so I am closing this issue. Thanks!