Skip to content

[MLIR][XeGPU] Add anchor_layout and update propagation to honor user-specified layouts #169267

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 9 commits into from
Nov 27, 2025
Merged

[MLIR][XeGPU] Add anchor_layout and update propagation to honor user-specified layouts #169267

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
b3f2a4a
adding anchor layout for load/store/prefetch_nd and dpas
Jianhui-Li Nov 22, 2025
bfae01f
propogation hornor pre-defined layout at anchor op
Jianhui-Li Nov 22, 2025
0482234
adding documentation
Jianhui-Li Nov 24, 2025
d1652af
address feedback and add more documentation
Jianhui-Li Nov 25, 2025
b186bc2
rename anchor_layout to layout
Jianhui-Li Nov 26, 2025
04f3edf
git merge main
Jianhui-Li Nov 26, 2025
60f5396
fix test
Jianhui-Li Nov 26, 2025
72fa240
fix clang-format
Jianhui-Li Nov 26, 2025
5f25c89
fix missing space in .clang-format
Jianhui-Li Nov 26, 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
223 changes: 173 additions & 50 deletions mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
View file
Open in desktop

Large diffs are not rendered by default.

2 changes: 2 additions & 0 deletions mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@ using std::optional;
namespace mlir {
namespace xegpu {

//#include "mlir/Dialect/XeGPU/IR/XeGPUOpInterface.cpp.inc"

void XeGPUDialect::initialize() {
addTypes<
#define GET_TYPEDEF_LIST
Expand Down
30 changes: 17 additions & 13 deletions mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -465,7 +465,7 @@ void PrefetchNdOp::build(OpBuilder &builder, OperationState &state,
xegpu::CachePolicyAttr l3_hint) {

return build(builder, state, tensorDesc, ValueRange(), DenseI64ArrayAttr(),
l1_hint, l2_hint, l3_hint);
l1_hint, l2_hint, l3_hint, /*anchor_layout=*/nullptr);
}

void PrefetchNdOp::build(OpBuilder &builder, OperationState &state,
Expand All @@ -480,7 +480,7 @@ void PrefetchNdOp::build(OpBuilder &builder, OperationState &state,
auto staticOffsetsAttr = builder.getDenseI64ArrayAttr(staticOffsets);

build(builder, state, tensorDesc, dynamicOffsets, staticOffsetsAttr, l1_hint,
l2_hint, l3_hint);
l2_hint, l3_hint, /*anchor_layout=*/nullptr);
}

LogicalResult PrefetchNdOp::verify() {
Expand Down Expand Up @@ -519,7 +519,7 @@ void LoadNdOp::build(OpBuilder &builder, OperationState &state, Type retType,

return build(builder, state, retType, tensorDesc, ValueRange(),
DenseI64ArrayAttr(), packed, transpose, l1_hint, l2_hint,
l3_hint);
l3_hint, /*anchor_layout=*/nullptr);
}

void LoadNdOp::build(OpBuilder &builder, OperationState &state, Type retType,
Expand All @@ -535,7 +535,8 @@ void LoadNdOp::build(OpBuilder &builder, OperationState &state, Type retType,
auto staticOffsetsAttr = builder.getDenseI64ArrayAttr(staticOffsets);

build(builder, state, retType, tensorDesc, dynamicOffsets, staticOffsetsAttr,
packed, transpose, l1_hint, l2_hint, l3_hint);
packed, transpose, l1_hint, l2_hint, l3_hint,
/*anchor_layout=*/nullptr);
}

LogicalResult LoadNdOp::verify() {
Expand Down Expand Up @@ -638,7 +639,8 @@ void StoreNdOp::build(OpBuilder &builder, OperationState &state, Value value,
xegpu::CachePolicyAttr l3_hint) {

return build(builder, state, value, tensorDesc, ValueRange(),
DenseI64ArrayAttr(), l1_hint, l2_hint, l3_hint);
DenseI64ArrayAttr(), l1_hint, l2_hint, l3_hint,
/*anchor_layout=*/nullptr);
}

void StoreNdOp::build(OpBuilder &builder, OperationState &state, Value value,
Expand All @@ -653,7 +655,7 @@ void StoreNdOp::build(OpBuilder &builder, OperationState &state, Value value,
auto staticOffsetsAttr = builder.getDenseI64ArrayAttr(staticOffsets);

build(builder, state, value, tensorDesc, dynamicOffsets, staticOffsetsAttr,
l1_hint, l2_hint, l3_hint);
l1_hint, l2_hint, l3_hint, /*anchor_layout=*/nullptr);
}

LogicalResult StoreNdOp::verify() {
Expand Down Expand Up @@ -826,7 +828,7 @@ void PrefetchOp::build(OpBuilder &builder, OperationState &state, Value source,
xegpu::CachePolicyAttr l2_hint,
xegpu::CachePolicyAttr l3_hint) {
build(builder, state, source, Value(), l1_hint, l2_hint, l3_hint,
IntegerAttr{});
IntegerAttr{}, /*anchor_layout=*/nullptr);
}

//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -876,7 +878,7 @@ void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
xegpu::CachePolicyAttr l2_hint,
xegpu::CachePolicyAttr l3_hint) {
build(builder, state, valueType, source, Value(), mask, IntegerAttr(),
l1_hint, l2_hint, l3_hint, /*layout=*/nullptr);
l1_hint, l2_hint, l3_hint, /*anchor_layout=*/nullptr);
}

void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
Expand All @@ -892,7 +894,7 @@ void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
auto offset = vector::FromElementsOp::create(builder, loc, type, values);

build(builder, state, valueType, source, offset, mask, chunk_size, l1_hint,
l2_hint, l3_hint, /*layout=*/nullptr);
l2_hint, l3_hint, /*anchor_layout=*/nullptr);
}

void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
Expand Down Expand Up @@ -960,7 +962,7 @@ void StoreScatterOp::build(OpBuilder &builder, OperationState &state,
xegpu::CachePolicyAttr l2_hint,
xegpu::CachePolicyAttr l3_hint) {
build(builder, state, value, dest, Value(), mask, IntegerAttr(), l1_hint,
l2_hint, l3_hint, /*layout=*/nullptr);
l2_hint, l3_hint, /*anchor_layout=*/nullptr);
}

void StoreScatterOp::build(OpBuilder &builder, OperationState &state,
Expand All @@ -978,7 +980,7 @@ void StoreScatterOp::build(OpBuilder &builder, OperationState &state,

// Call the correct builder overload that does not expect result types.
build(builder, state, value, dest, offset, mask, chunk_size, l1_hint, l2_hint,
l3_hint, /*layout=*/nullptr);
l3_hint, /*anchor_layout=*/nullptr);
}

void StoreScatterOp::build(
Expand Down Expand Up @@ -1155,7 +1157,8 @@ LogicalResult LoadMatrixOp::verify() {
MemDescType mdescTy = getMemDesc().getType();

return IsValidMatrixOpParams(resTy, mdescTy, subgroup_block_io,
getLayoutAttr(), [&]() { return emitError(); });
getAnchorLayoutAttr(),
[&]() { return emitError(); });
}

//===----------------------------------------------------------------------===//
Expand All @@ -1179,7 +1182,8 @@ LogicalResult StoreMatrixOp::verify() {
UnitAttr subgroup_block_io = getSubgroupBlockIoAttr();
MemDescType mdescTy = getMemDesc().getType();
return IsValidMatrixOpParams(dataTy, mdescTy, subgroup_block_io,
getLayoutAttr(), [&]() { return emitError(); });
getAnchorLayoutAttr(),
[&]() { return emitError(); });
}

namespace mlir {
Expand Down
Loading