[LinalgToXeGPU] Lower linalg.matmul_transpose_b into xegpu.dpas

[dchigarev]

Closes #340

Support lowering of linalg.matmul_transpose_b as well as linalg.transpose %b + linalg.matmul %a %b into xegpu.dpas.

To make a transposed multiplication we have to load B chunks with xegpu.load_nd ... <transpose = array<i64: 1, 0>> attribute and also change the iteration dimension (from rows to cols).

[dchigarev]

changes iteration dimension for B chunks

[kurapov-peter]

This seem to be almost generic enough to support transpose A and the batched transposed versions. Could you try making the interfaces amenable to all the matmul variations? Implementation can stay as is.

[kurapov-peter]

What would alloca mean?

[dchigarev]

nothing, it was added here by mistake, will remove in the next commit

[dchigarev]

This seem to be almost generic enough to support transpose A and the batched transposed versions. Could you try making the interfaces amenable to all the matmul variations? Implementation can stay as is.

1. Regarding batched transpose: at the moment there's no support for either linalg.batch_matmul or for linalg.batch_reduce_matmul in linalg-to-xegpu pass. Trying YOLO approach by simply adding those ops to the list of allowed ops causes a segfault, so probably a different handling is needed in those cases. Anyway, the 'transpose_b' logic shouldn't change at all if createDPASKernel() would support batched matmuls sometime. You would just need to add BatchMatmulOp everywhere and it should work.

2. Regarding 'transpose A': this one also requires special handling. Following the current logic of createDPASKernel() we can't use transpose attribute of xegpu.load_nd to transpose A argument since it only can transpose 32bit blocks (and we only support f16 in our lowering). This means, that instead of using xegpu.load_nd <transpose> we would need to use vector.transpose on the loaded vector. Since it doesn't correlate with the changes in this PR and seems like a separate feature, I would propose to consider it as a separate issue ([LinalgToXeGPU] Support linalg.matmul_transpose_a #361). Anyway, the general logic of detecting "if there's a transpose before matmul" stays the same. Added an argument to findAndReplaceTranspose() function indicating which matmul operand we're processing.

   Example of why 32bit transpose causes problems for f16 data

   A(f16): // original matrix
   [1,   2,  3,  4]
   [5,   6,  7,  8]
   [9,  10, 11, 12]
   [13, 14, 15, 16]
   
   xegpu.load_nd A(f16) <transpose>: // transposed 32bit blocks
   [1, 2,  9, 10]
   [5, 6, 13, 14]
   [3, 4, 11, 12]
   [7, 8, 15, 16]
   
   vector.transpose A(f16): // expected transpose result
   [1, 5,  9, 13]
   [2, 6, 10, 14]
   [3, 7, 11, 15]
   [4, 8, 12, 16]
   

   Why it doesn't cause any problems for B argument

   In contrast with 'A' argument we load 'B' with 'packed/vnni' attribute which apparently involves some special handling which is properly aligned with our needs of B transpose.

   P.S. there's transpose_bit_width attribute that should in theory solve this problem, but setting it to 16 does nothing. Not sure if this an intended behavior or just a bug in xegpu-to-vc-func lowering. Submitted a question to IMEX ([xegpu-to-vc-func] Is transpose_bit_width=16 supported? mlir-extensions#895).

[kurapov-peter]

This seem to be almost generic enough to support transpose A and the batched transposed versions. Could you try making the interfaces amenable to all the matmul variations? Implementation can stay as is.

1. Regarding batched transpose: at the moment there's no support for either linalg.batch_matmul or for linalg.batch_reduce_matmul in linalg-to-xegpu pass. Trying YOLO approach by simply adding those ops to the list of allowed ops causes a segfault, so probably a different handling is needed in those cases. Anyway, the 'transpose_b' logic shouldn't change at all if createDPASKernel() would support batched matmuls sometime. You would just need to add BatchMatmulOp everywhere and it should work.

2. Regarding 'transpose A': this one also requires special handling. Following the current logic of createDPASKernel() we can't use transpose attribute of xegpu.load_nd to transpose A argument since it only can transpose 32bit blocks (and we only support f16 in our lowering). This means, that instead of using xegpu.load_nd <transpose> we would need to use vector.transpose on the loaded vector. Since it doesn't correlate with the changes in this PR and seems like a separate feature, I would propose to consider it as a separate issue ([LinalgToXeGPU] Support linalg.matmul_transpose_a #361). Anyway, the general logic of detecting "if there's a transpose before matmul" stays the same. Added an argument to findAndReplaceTranspose() function indicating which matmul operand we're processing.
   Example of why 32bit transpose causes problems for f16 data

   A(f16): // original matrix
   [1,   2,  3,  4]
   [5,   6,  7,  8]
   [9,  10, 11, 12]
   [13, 14, 15, 16]
   
   xegpu.load_nd A(f16) <transpose>: // transposed 32bit blocks
   [1, 2,  9, 10]
   [5, 6, 13, 14]
   [3, 4, 11, 12]
   [7, 8, 15, 16]
   
   vector.transpose A(f16): // expected transpose result
   [1, 5,  9, 13]
   [2, 6, 10, 14]
   [3, 7, 11, 15]
   [4, 8, 12, 16]
   

   Why it doesn't cause any problems for B argument
   In contrast with 'A' argument we load 'B' with 'packed/vnni' attribute which apparently involves some special handling which is properly aligned with our needs of B transpose.

   P.S. there's transpose_bit_width attribute that should in theory solve this problem, but setting it to 16 does nothing. Not sure if this an intended behavior or just a bug in xegpu-to-vc-func lowering. Submitted a question to IMEX ([xegpu-to-vc-func] Is transpose_bit_width=16 supported? mlir-extensions#895).

Sure, I did not suggest adding support for other flavors right away, only to make the interfaces friendly. Supporting other matmuls is a separate issue.

[kurapov-peter]

There's also the bool transpose param in a couple places that could be generalized but that's a minor thing.

[dchigarev]

There's also the bool transpose param in a couple places that could be generalized but that's a minor thing.

This bool parameter is used in functions that also take a matmul operand, so this parameter can be used with both of the operands (and with any matmul op). Is this generic enough or you had something else on your mind?

bool transposeA = true;
bool transposeB = false;

auto tilesA = createTiles(matA, ..., /*transpose=*/transposeA);
auto tilesB = createTiles(matB, ..., /*transpose=*/transposeB);

[kurapov-peter]

There's also the bool transpose param in a couple places that could be generalized but that's a minor thing.

This bool parameter is used in functions that also take a matmul operand, so this parameter can be used with both of the operands (and with any matmul op). Is this generic enough or you had something else on your mind?

bool transposeA = true;

bool transposeB = false;



auto tilesA = createTiles(matA, ..., /*transpose=*/transposeA);

auto tilesB = createTiles(matB, ..., /*transpose=*/transposeB);

Yup. Sounds good!