Skip to content

Conversation

@dchigarev
Copy link
Contributor

@dchigarev dchigarev commented Nov 29, 2024

This PR allows any-D memrefs to be used in linalg-to-xegpu pass as long they can be squeezed to a 2D memref through leading dimensions. Example:

// allowed
!chunk_type = memref<1x1x8x16xf16, strided<[512, 128, 16, 1], offset: ?>>
linalg.mul ins(%0, %1 : !chunk_type, !chunk_type) outs(%2 : !chunk_type)

// NOT allowed
!chunk_type = memref<1x8x1x16xf16, strided<[512, 128, 16, 1], offset: ?>>
linalg.mul ins(%0, %1 : !chunk_type, !chunk_type) outs(%2 : !slm_chunk)

The linalg-to-xegpu logic considers a linalg operation that can be lowered to xegpu and checks the operand shapes. The old logic refused all linalgOps that have non-2D operands. The new logic additionally checks if a non-2D operand can be squeezed to a 2D and allows them if they can. It then squeezes the non-2D operand using memref.subview and replaces old operands with new 2D memrefs. Example:

// input:
!chunk_type = memref<1x1x8x16xf16, strided<[512, 128, 16, 1], offset: ?>>
linalg.mul ins(%0, %1 : !chunk_type, !chunk_type) outs(%2 : !chunk_type)

// output:
!chunk_type = memref<1x1x8x16xf16, strided<[512, 128, 16, 1], offset: ?>>
!squeezed_type = memref<8x16xf16, strided<[16, 1], offset: ?>>

%0_sq = memref.subview %0[0, 0, 0, 0] [1, 1, 8, 16] [1, 1, 1, 1] : !chunk_type to !squeezed_type
%1_sq = memref.subview %1[0, 0, 0, 0] [1, 1, 8, 16] [1, 1, 1, 1] : !chunk_type to !squeezed_type
%2_sq = memref.subview %2[0, 0, 0, 0] [1, 1, 8, 16] [1, 1, 1, 1] : !chunk_type to !squeezed_type
linalg.mul ins(%0_sq, %1_sq : !squeezed_type , !squeezed_type ) outs(%2_sq : !squeezed_type )

Comment on lines 208 to 217
while (auto subViewOp = memref.getDefiningOp<memref::SubViewOp>()) {
auto currentOffsets = getAsOpFoldResult(resolvedOffsets);
resolvedOffsets.clear();

affine::resolveIndicesIntoOpWithOffsetsAndStrides(
rewriter, memref.getLoc(), subViewOp.getMixedOffsets(),
subViewOp.getMixedStrides(), subViewOp.getDroppedDims(), currentOffsets,
resolvedOffsets);
memref = subViewOp.getOperand(0);
}
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the logic was taken from XeGPUFoldAliasOps upstream pass

@dchigarev dchigarev marked this pull request as ready for review November 29, 2024 15:20
Signed-off-by: dchigarev <[email protected]>
Signed-off-by: dchigarev <[email protected]>
Copy link
Contributor

@kurapov-peter kurapov-peter left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some minor comments, fine otherwise

@dchigarev dchigarev merged commit 472ef9c into intel:main Dec 3, 2024
6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants