Page MenuHomePhabricator

Add a pattern to convert static linalg.matmul -> linalg.mmt4d

Authored by asaadaldien on Jul 14 2021, 12:33 PM.



This pattern converts statically shaped linalg.matmul when operands sizes (m, n, k) are an integer multiple of
mmt4d's operands inner most dims (m0, n0, k0).

This pattern does the following transofmrations to operands and mmt4d results:


lhs: (m, k) -(reshape)-> (m1, m0, k1, k0) -(transpose)-> (m1, k1, m0, k0)
rhs: (k, n) -(reshape)-> (k1, k0, n1, n0) -(transpose)-> (n1, k1, n0, k0)


mmt4d: (m1, n1, m0, n0) -(transpose)-> (m1, m0, n1, n0) -(reshape)-> (m, n)

Diff Detail

Event Timeline

asaadaldien created this revision.Jul 14 2021, 12:33 PM
asaadaldien requested review of this revision.Jul 14 2021, 12:33 PM
asaadaldien edited the summary of this revision. (Show Details)
Benoit added inline comments.Jul 14 2021, 12:45 PM

Add a comment that it's temporary that we are bailing out here, and that we will eventually want to accept any sizes by doing padding?


I'm thinking that it might help readability of this code if this function were named 'packMatrix' : "transpose" is an implementation detail, and the concept of "transposition" is a bit overloaded in the present context since we are also dealing with transposing the RHS matrix (which gives the letter T in the name, "MMT4D").

'Matrix' vs 'Operand' because I'm thinking, this function in itself only cares that it gets a matrix to work on; that matrix only needs to be called an 'operand' in the context of the mmt4d op.


If the above function were renamed to packMatrix, then it might make sense to rename this one to unpackMatrix.

Benoit added inline comments.Jul 14 2021, 12:48 PM

Ah sorry I had missed that the 2D->4D conversion was done above separately in expandTo4D.

Do you think that it would be reasonable to call expandTo4D inside packMatrix (and unpackMatrix below) i.e. my proposed new names for transposeOperand, collapseOperand, so that the internals of packing/unpacking are abstracted behind these functions and the higher-level code here can thus look simpler?

asaadaldien added inline comments.Jul 14 2021, 1:41 PM

We have to check divisibility here as a precondition, padding is something that is handled outside this pattern (e.g we do pad operands in IREE to an next integer multiple prior to apply this).

Benoit added inline comments.Jul 14 2021, 1:45 PM

I see; I understand how in the current context where K0,M0,N0 are compile-time constants from the get go (ie already here where we are lowering matmul to mmt4d) , it's possible to split the padding part into a separate pattern.

But how is this going to generalize to when K0,M0,N0 are '?' dynamic shape dimensions at the time when we need to convert matmul to mmt4d?

asaadaldien added inline comments.Jul 14 2021, 2:00 PM

There will be another pattern that doesn't have M0, N0, K0 as constants and doesn't have to run on statically shaped operands as it will do <?x?xf32> -> <?x?x?x?xf32> which linalg.tensor_expand_shape doesn't support atm. (In addition to compiler support to specialize M0, N0, K0 for runtime versioning in IREE...etc).

asaadaldien abandoned this revision.Aug 2 2021, 1:29 PM

We are adding this pattern to IREE for now..