Skip to content

[mlir][XeGPU] add unroll patterns for load_matrix and store_matrix #154637

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Sep 3, 2025
Merged

[mlir][XeGPU] add unroll patterns for load_matrix and store_matrix #154637

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
22d4193
add unroll pattern and unit test for load_matrix and store_matrix
chencha3 Aug 20, 2025
d20da85
add unroll pattern and unit test for load_matrix and store_matrix
chencha3 Aug 20, 2025
87e4fd2
Merge branch 'ld_st_matrix_unroll_patterns' of https://github.com/che…
chencha3 Aug 21, 2025
b3af260
Merge branch 'main' into ld_st_matrix_unroll_patterns
chencha3 Aug 27, 2025
442c18a
merge
chencha3 Aug 27, 2025
a368430
roll back unnecessary change
chencha3 Aug 27, 2025
3f5d692
add unit test
chencha3 Aug 27, 2025
a0512d9
address comments
chencha3 Sep 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,8 +67,8 @@ def XeGPUBlocking: Pass<"xegpu-blocking"> {
to a hardware instruction.
}];
let dependentDialects = [
"memref::MemRefDialect", "xegpu::XeGPUDialect", "vector::VectorDialect"
];
"memref::MemRefDialect", "xegpu::XeGPUDialect", "vector::VectorDialect",
"index::IndexDialect"];
Copy link
Contributor

Choose a reason for hiding this comment

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

nit: does it use index dialect at all in the end?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Yes, it is used in addElementwise, so the pass needs to load it.

}

#endif // MLIR_DIALECT_XEGPU_TRANSFORMS_PASSES_TD
1 change: 1 addition & 0 deletions mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ class OpResult;
class OpBuilder;
class ValueRange;
class TypeConverter;
class OpFoldResult;

namespace xegpu {
class DistributeLayoutAttr;
Expand Down
5 changes: 3 additions & 2 deletions mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

#include "mlir/Dialect/XeGPU/Transforms/Passes.h"

#include "mlir/Dialect/Index/IR/IndexDialect.h"
#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
Expand Down Expand Up @@ -157,10 +158,10 @@ XeGPUBlockingPass::getTileShape(const T &operandOrResult) const {
std::optional<SmallVector<int64_t>>
XeGPUBlockingPass::getTileShape(Operation *op) const {
if (isa<xegpu::CreateNdDescOp, xegpu::UpdateNdOffsetOp, xegpu::CreateDescOp,
xegpu::UpdateOffsetOp>(op))
xegpu::UpdateOffsetOp, xegpu::LoadMatrixOp>(op))
return getTileShape(op->getOpResult(0));
if (isa<xegpu::PrefetchNdOp, xegpu::LoadNdOp, xegpu::PrefetchOp,
xegpu::LoadGatherOp>(op))
xegpu::LoadGatherOp, xegpu::StoreMatrixOp>(op))
return getTileShape(op->getOpOperand(0));
if (isa<xegpu::StoreNdOp, xegpu::StoreScatterOp>(op))
return getTileShape(op->getOpOperand(1));
Expand Down
87 changes: 82 additions & 5 deletions mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -682,13 +682,90 @@ struct UnrollUpdateOffsetOp : public UnrollPattern<xegpu::UpdateOffsetOp> {
}
};

struct UnrollLoadMatrixOp : public UnrollPattern<xegpu::LoadMatrixOp> {
using UnrollPattern<xegpu::LoadMatrixOp>::UnrollPattern;
LogicalResult matchAndRewrite(xegpu::LoadMatrixOp op,
PatternRewriter &rewriter) const override {
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape)
return failure();

Location loc = op.getLoc();
VectorType valueTy = op.getType();
Type elemTy = valueTy.getElementType();
ArrayRef<int64_t> shape = valueTy.getShape();
auto layout = dyn_cast<xegpu::LayoutAttr>(op.getLayoutAttr());

VectorType newValueTy = valueTy.cloneWith(*targetShape, elemTy);

SmallVector<OpFoldResult> mixedOffsets = op.getMixedOffsets();
SmallVector<SmallVector<OpFoldResult>> offsetsList;
for (SmallVector<int64_t> offsets :
StaticTileOffsetRange(shape, *targetShape)) {
auto adds = xegpu::addWithRightAligned(
rewriter, loc, mixedOffsets,
getAsIndexOpFoldResult(op.getContext(), offsets));
offsetsList.push_back(adds);
}

SmallVector<Value> newOps;
for (SmallVector<OpFoldResult> offsets : offsetsList) {
auto newOp = rewriter.create<xegpu::LoadMatrixOp>(
op.getLoc(), newValueTy, op.getMemDesc(), offsets,
layout.dropInstData());
newOps.push_back(newOp);
}
Value castOp = unpack(newOps, op.getType(), *targetShape, loc, rewriter);
rewriter.replaceOp(op, castOp);
return success();
}
};

struct UnrollStoreMatrixOp : public UnrollPattern<xegpu::StoreMatrixOp> {
using UnrollPattern<xegpu::StoreMatrixOp>::UnrollPattern;
LogicalResult matchAndRewrite(xegpu::StoreMatrixOp op,
PatternRewriter &rewriter) const override {
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape)
return failure();

Location loc = op.getLoc();
VectorType valueTy = op.getData().getType();
ArrayRef<int64_t> shape = valueTy.getShape();
auto layout = dyn_cast<xegpu::LayoutAttr>(op.getLayoutAttr());

SmallVector<Type> convertedValTypes =
getUnrolledTypes(valueTy, *targetShape);
SmallVector<Value> convertedValues =
pack(op.getData(), convertedValTypes, *targetShape, loc, rewriter);

SmallVector<OpFoldResult> mixedOffsets = op.getMixedOffsets();
SmallVector<SmallVector<OpFoldResult>> offsetsList;
for (SmallVector<int64_t> offsets :
StaticTileOffsetRange(shape, *targetShape)) {
auto adds = xegpu::addWithRightAligned(
Copy link
Contributor

Choose a reason for hiding this comment

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

I don't think we need to use addWithRightAligned here? The op's offsets should have always the same number as the distributed offsets (out from shape/targetshape).

Copy link
Contributor Author

Choose a reason for hiding this comment

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

This is motivated by the old style (offset in createNd), in which creating a low rank tensor desc (2D) from a high rank tensor desc (e.g., 4D) was allowed. In such case, the local offset for the tensor desc is with lower rank, which needs to be added with the original high rank offsets to get the final one. It is created to make it compatible with these test cases. I am also not sure whether this support will be removed or not. It can be refactored after we completely switch to the new style.

Copy link
Contributor

Choose a reason for hiding this comment

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

I don't think this applies to load_matrix/store_matrix: we should restrict the shapes size are always 2D - for both matrix_desc and vector.

Copy link
Contributor

Choose a reason for hiding this comment

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

AFAIK, it's already guaranteed to be 2D by matrix ops verifiers.
It should be fine as is but if you can use a different helper/approach for better readability, that'd be nice too.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

addWithRightAligned was supposed to handle both cases. But I think your suggestions are right. So I introduced addElementwise here.

rewriter, loc, mixedOffsets,
getAsIndexOpFoldResult(op.getContext(), offsets));
offsetsList.push_back(adds);
}

for (auto [v, offsets] : llvm::zip_equal(convertedValues, offsetsList))
rewriter.create<xegpu::StoreMatrixOp>(loc, v, op.getMemDesc(), offsets,
layout.dropInstData());

rewriter.eraseOp(op);
return success();
}
};

} // namespace

void mlir::xegpu::populateXeGPUUnrollPatterns(
RewritePatternSet &patterns, const xegpu::UnrollOptions &options) {
patterns.add<UnrollCreateNdOp, UnrollUpdateNdOffsetOp, UnrollPrefetchNdOp,
UnrollLoadNdOp, UnrollStoreNdOp, UnrollDpasOp,
UnrollCreateDescOp, UnrollLoadGatherOp, UnrollStoreScatterOp,
UnrollPrefetchOp, UnrollUpdateOffsetOp>(patterns.getContext(),
options);
patterns
.add<UnrollCreateNdOp, UnrollUpdateNdOffsetOp, UnrollPrefetchNdOp,
UnrollLoadNdOp, UnrollStoreNdOp, UnrollDpasOp, UnrollCreateDescOp,
UnrollLoadGatherOp, UnrollStoreScatterOp, UnrollPrefetchOp,
UnrollUpdateOffsetOp, UnrollLoadMatrixOp, UnrollStoreMatrixOp>(
patterns.getContext(), options);
}
19 changes: 18 additions & 1 deletion mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,14 @@ xegpu::DistributeLayoutAttr xegpu::getDistributeLayoutAttr(const Value value) {
if (auto loadNd = dyn_cast<xegpu::LoadNdOp>(defOp))
return getDistributeLayoutAttr(loadNd.getTensorDesc());

// for LoadMatrixOp, the layout is attached to the property of the op
if (auto loadOp = dyn_cast<xegpu::LoadMatrixOp>(defOp))
return loadOp.getLayoutAttr();

// for StoreMatrixOp, the layout is attached to the property of the op
if (auto storeOp = dyn_cast<xegpu::StoreMatrixOp>(defOp))
return storeOp.getLayoutAttr();

std::string layoutName = getLayoutName(result);
if (defOp->hasAttr(layoutName))
return defOp->getAttrOfType<xegpu::DistributeLayoutAttr>(layoutName);
Expand All @@ -154,6 +162,13 @@ xegpu::DistributeLayoutAttr xegpu::getDistributeLayoutAttr(const Value value) {
xegpu::DistributeLayoutAttr
xegpu::getDistributeLayoutAttr(const OpOperand &opr) {
Operation *op = opr.getOwner();

if (auto loadOp = dyn_cast<xegpu::LoadMatrixOp>(op))
return loadOp.getLayoutAttr();

if (auto storeOp = dyn_cast<xegpu::StoreMatrixOp>(op))
return storeOp.getLayoutAttr();

std::string layoutName = xegpu::getLayoutName(opr);
if (op->hasAttr(layoutName))
return op->getAttrOfType<xegpu::DistributeLayoutAttr>(layoutName);
Expand Down Expand Up @@ -182,6 +197,9 @@ template void xegpu::setDistributeLayoutAttr<mlir::OpOperand>(
void xegpu::setDistributeLayoutAttrs(
Operation *op, function_ref<DistributeLayoutAttr(Value)> getLayoutImpl) {
op->walk([&](Operation *nestOp) {
if (isa<xegpu::LoadMatrixOp, xegpu::StoreMatrixOp>(nestOp))
return;

for (OpOperand &opr : nestOp->getOpOperands()) {
auto layout = getLayoutImpl(opr.get());
setDistributeLayoutAttr(opr, layout);
Expand Down Expand Up @@ -454,5 +472,4 @@ xegpu::addWithRightAligned(OpBuilder &builder, Location loc,
results.push_back(builder.createOrFold<index::AddOp>(loc, lval, rval));
}
return results;
return {};
}
23 changes: 23 additions & 0 deletions mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -561,3 +561,26 @@ gpu.module @test_kernel {
gpu.return %e : vector<8x32x2xf16>
}
}

// -----
gpu.module @test_kernel {
//CHECK-LABEL: unroll_load_matrix
gpu.func @unroll_load_matrix(%arg0: memref<4096xi8, 3>) -> vector<32x32xf32> {
%0 = xegpu.create_mem_desc %arg0 : memref<4096xi8, 3> -> !xegpu.mem_desc<32x32xf32>
//CHECK-COUNT-8: xegpu.load_matrix {{.*}} : !xegpu.mem_desc<32x32xf32>, index, index -> vector<8x16xf32>
//CHECK-COUNT-8: vector.insert_strided_slice {{.*}} : vector<8x16xf32> into vector<32x32xf32>
%1 = xegpu.load_matrix %0[0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}>: !xegpu.mem_desc<32x32xf32> -> vector<32x32xf32>
gpu.return %1: vector<32x32xf32>
}
}

// -----
gpu.module @test_kernel {
// CHECK-LABEL: unroll_store_matrix
gpu.func @unroll_store_matrix(%value: vector<32x32xf32>, %arg0 : memref<32768xi8, 3>) {
%mdesc = xegpu.create_mem_desc %arg0 : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>
// CHECK-COUNT-8: xegpu.store_matrix {{.*}} : vector<8x16xf32>, !xegpu.mem_desc<64x128xf32>, index, index
xegpu.store_matrix %value, %mdesc[0, 0] {layout = #xegpu.layout<inst_data = [8, 16]>} : vector<32x32xf32>, !xegpu.mem_desc<64x128xf32>
gpu.return
}
}
27 changes: 27 additions & 0 deletions mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,33 @@ gpu.module @test_1_1_assignment {
gpu.return
}

// CHECK-LABEL: create_nd_tdesc_from_higher_rank_memref
// CHECK-SAME: [[ARG_0:%.*]]: memref<3x256x128xf32>
gpu.func @create_nd_tdesc_from_higher_rank_memref(%src: memref<3x256x128xf32>) {
//CHECK: [[SGID:%.+]] = gpu.subgroup_id : index
//CHECK: [[SGIDY:%.+]] = affine.apply #map()[[[SGID]]]
//CHECK: [[SGIDX:%.+]] = affine.apply #map1()[[[SGID]]]
//CHECK: [[C32:%.+]] = arith.constant 32 : index
//CHECK: [[LY:%.+]] = index.mul [[SGIDY]], [[C32]]
//CHECK: [[LX:%.+]] = index.mul [[SGIDX]], [[C32]]
//CHECK: [[C0:%.+]] = arith.constant 0 : index
//CHECK: [[C0_2:%.+]] = arith.constant 0 : index
//CHECK: [[UY:%.+]] = arith.addi [[LY]], [[C0]] : index
//CHECK: [[UX:%.+]] = arith.addi [[LX]], [[C0_2]] : index
//CHECK: [[C256:%.+]] = arith.constant 256 : index
//CHECK: [[MODY:%.+]] = index.remu [[UY]], [[C256]]
//CHECK: [[C128:%.+]] = arith.constant 128 : index
//CHECK: [[MODX:%.+]] = index.remu [[UX]], [[C128]]
//CHECK: [[C0_3:%.+]] = arith.constant 0 : index
//CHECK: [[Y:%.+]] = index.add [[MODY]], [[C0_3]]
//CHECK: [[C0_4:%.+]] = arith.constant 0 : index
//CHECK: [[X:%.+]] = index.add [[MODX]], [[C0_4]]
//CHECK: [[TDESC:%.+]] = xegpu.create_nd_tdesc [[ARG_0]][1, [[Y]], [[X]]] : memref<3x256x128xf32> -> !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
%tdesc = xegpu.create_nd_tdesc %src[1, 0, 0] : memref<3x256x128xf32>
-> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>
gpu.return
}

// CHECK-LABEL: load_nd_tdesc
// CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>
gpu.func @load_nd_tdesc(%src: memref<256x128xf32>) {
Expand Down