[Mlir-commits] [mlir] [MLIR][XeGPU][TransformOps] set_op_layout_attr supports setting anchor layout (PR #172542)
Tuomas Kärnä
llvmlistbot at llvm.org
Thu Feb 12 09:03:46 PST 2026
https://github.com/tkarna updated https://github.com/llvm/llvm-project/pull/172542
>From 75e175890afbd7f298297b738250af504080c8bd Mon Sep 17 00:00:00 2001
From: Tuomas Karna <tuomas.karna at intel.com>
Date: Fri, 12 Dec 2025 19:31:01 +0200
Subject: [PATCH 1/4] set_op_layout_attr supports setting anchor layout
---
.../XeGPU/TransformOps/XeGPUTransformOps.td | 22 ++++++++-----
.../XeGPU/TransformOps/XeGPUTransformOps.cpp | 33 ++++++++++++++++---
mlir/python/mlir/dialects/transform/xegpu.py | 4 +++
.../Dialect/XeGPU/transform-ops-invalid.mlir | 21 +++++++++++-
mlir/test/Dialect/XeGPU/transform-ops.mlir | 24 ++++++++++++--
.../python/dialects/transform_xegpu_ext.py | 32 ++++++++++++++++--
6 files changed, 117 insertions(+), 19 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td b/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
index 29579acc727ed..8c5539ad19382 100644
--- a/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
@@ -109,12 +109,14 @@ def SetOpLayoutAttrOp : Op<Transform_Dialect, "xegpu.set_op_layout_attr", [
let summary = "Set xegpu.layout attribute of an op.";
let description = [{
- Sets the `xegpu.layout` attribute of an op. If `result=true`, sets the
- `layout_result_{index}`, otherwise `layout_operand_{index}` attribute. The
- target operand/result value is defined by the `index` argument. The layout
- is defined by the `sg_layout`, `sg_data` and optional `inst_data`
- attributes. If `slice_dims` is provided, the `xegpu.layout` attribute is
- wrapped in an `xegpu.slice<..., dims=slice_dims>` attribute.
+ Sets the `xegpu.layout` attribute of an op. By default it sets the anchor
+ layout for XeGPU ops that support it. If `result=true` or `operand=true`,
+ it sets the `layout_result_{index}` or `layout_operand_{index}` attribute,
+ respectively, applicable to any op. The target operand/result value is
+ defined by the `index` argument. The layout is defined by the `sg_layout`,
+ `sg_data` and optional `inst_data` attributes. If `slice_dims` is provided,
+ the `xegpu.layout` attribute is wrapped in an `xegpu.slice<..., dims=slice_dims>`
+ attribute.
}];
let arguments = (ins TransformHandleTypeInterface:$target,
@@ -126,7 +128,8 @@ def SetOpLayoutAttrOp : Op<Transform_Dialect, "xegpu.set_op_layout_attr", [
DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_sg_data,
DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_inst_data,
DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$slice_dims,
- DefaultValuedAttr<UnitAttr, "false">:$result
+ DefaultValuedAttr<UnitAttr, "false">:$result,
+ DefaultValuedAttr<UnitAttr, "false">:$operand
);
let results = (outs);
@@ -137,12 +140,13 @@ def SetOpLayoutAttrOp : Op<Transform_Dialect, "xegpu.set_op_layout_attr", [
"ArrayRef<OpFoldResult>":$mixedSgData,
"ArrayRef<OpFoldResult>":$mixedInstData,
"ArrayRef<int64_t>":$sliceDims,
- CArg<"bool", "false">:$result
+ CArg<"bool", "false">:$result,
+ CArg<"bool", "false">:$operand
)>,
];
let assemblyFormat = [{
- $target (`result` $result^)? (`index` `=` $index^)?
+ $target (`result` $result^)? (`operand` $operand^)? (`index` `=` $index^)?
`sg_layout` `=` custom<DynamicIndexList>($sg_layout, $static_sg_layout)
`sg_data` `=` custom<DynamicIndexList>($sg_data, $static_sg_data)
(`inst_data` `=` custom<DynamicIndexList>($inst_data, $static_inst_data)^)?
diff --git a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
index 285dfac56be08..e4bcff4b996eb 100644
--- a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
@@ -291,7 +291,7 @@ void transform::SetOpLayoutAttrOp::build(
OpBuilder &builder, OperationState &ostate, Value target, int64_t index,
ArrayRef<OpFoldResult> mixedSgLayout, ArrayRef<OpFoldResult> mixedSgData,
ArrayRef<OpFoldResult> mixedInstData, ArrayRef<int64_t> sliceDims,
- bool result) {
+ bool result, bool operand) {
SmallVector<int64_t> staticSgLayout, staticSgData, staticInstData;
SmallVector<Value> dynamicSgLayout, dynamicSgData, dynamicInstData;
dispatchIndexOpFoldResults(mixedSgLayout, dynamicSgLayout, staticSgLayout);
@@ -307,7 +307,8 @@ void transform::SetOpLayoutAttrOp::build(
/*static_sg_data=*/staticSgData,
/*static_inst_data=*/staticInstData,
/*slice_dims=*/sliceDims,
- /*result=*/result);
+ /*result=*/result,
+ /*operand=*/operand);
}
DiagnosedSilenceableFailure
@@ -322,13 +323,19 @@ transform::SetOpLayoutAttrOp::apply(transform::TransformRewriter &rewriter,
Operation *target = *targetOps.begin();
bool resultTarget = getResult();
+ bool operandTarget = getOperand();
+
+ if (resultTarget && operandTarget) {
+ return emitSilenceableFailure(getLoc())
+ << "Cannot set both result and operand layout attributes.";
+ }
int64_t index = getIndex();
if (resultTarget && index >= target->getNumResults()) {
return emitSilenceableFailure(getLoc())
<< "Index exceeds the number of op results";
}
- if (!resultTarget && index >= target->getNumOperands()) {
+ if (operandTarget && index >= target->getNumOperands()) {
return emitSilenceableFailure(getLoc())
<< "Index exceeds the number of op operands";
}
@@ -349,10 +356,26 @@ transform::SetOpLayoutAttrOp::apply(transform::TransformRewriter &rewriter,
}
// Set layout attribute for the op result or operand
- if (resultTarget)
+ if (resultTarget) {
xegpu::setDistributeLayoutAttr(target->getResult(index), layout);
- else
+ } else if (operandTarget) {
xegpu::setDistributeLayoutAttr(target->getOpOperand(index), layout);
+ } else {
+ // Set anchor layout if requested.
+ // TODO use AnchorLayoutInterface when available.
+ if (!isa<xegpu::LoadNdOp>(target)) {
+ auto diag = emitSilenceableFailure(getLoc())
+ << "Cannot set anchor layout to op: " << target->getName();
+ diag.attachNote(target->getLoc()) << "target op";
+ return diag;
+ }
+ auto loadOp = dyn_cast<xegpu::LoadNdOp>(target);
+ if (loadOp)
+ loadOp.setLayoutAttr(layout);
+ auto prefetchOp = dyn_cast<xegpu::PrefetchOp>(target);
+ if (prefetchOp)
+ prefetchOp.setLayoutAttr(layout);
+ }
return DiagnosedSilenceableFailure::success();
}
diff --git a/mlir/python/mlir/dialects/transform/xegpu.py b/mlir/python/mlir/dialects/transform/xegpu.py
index 5aa6453b7cb8a..a768ce5f4e720 100644
--- a/mlir/python/mlir/dialects/transform/xegpu.py
+++ b/mlir/python/mlir/dialects/transform/xegpu.py
@@ -134,6 +134,7 @@ def __init__(
slice_dims: Optional[MixedInt] = None,
index: Optional[Union[int, Attribute]] = None,
result: Optional[Union[bool, Attribute]] = None,
+ operand: Optional[Union[bool, Attribute]] = None,
loc=None,
ip=None,
):
@@ -164,6 +165,7 @@ def __init__(
slice_dims=slice_dims,
index=index,
result=result,
+ operand=operand,
loc=loc,
ip=ip,
)
@@ -178,6 +180,7 @@ def set_op_layout_attr(
slice_dims: Optional[MixedInt] = None,
index: Optional[Union[int, Attribute]] = None,
result: Optional[Union[bool, Attribute]] = None,
+ operand: Optional[Union[bool, Attribute]] = None,
loc=None,
ip=None,
) -> SetOpLayoutAttrOp:
@@ -189,6 +192,7 @@ def set_op_layout_attr(
slice_dims=slice_dims,
index=index,
result=result,
+ operand=operand,
loc=loc,
ip=ip,
)
diff --git a/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir b/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir
index dce4a41982550..2a147497a893b 100644
--- a/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir
@@ -47,7 +47,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%0 = transform.structured.match ops{["arith.extf"]} in %arg1 : (!transform.any_op) -> !transform.any_op
// expected-error at below {{Index exceeds the number of op operands}}
- transform.xegpu.set_op_layout_attr %0 index = 1 sg_layout = [8, 4] sg_data = [32, 64] : !transform.any_op
+ transform.xegpu.set_op_layout_attr %0 operand index = 1 sg_layout = [8, 4] sg_data = [32, 64] : !transform.any_op
transform.yield
}
}
@@ -67,6 +67,25 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%0 = transform.structured.match ops{["arith.extf"]} in %arg1 : (!transform.any_op) -> !transform.any_op
// expected-error at below {{Requires exactly one targetOp handle (got 2)}}
+ transform.xegpu.set_op_layout_attr %0 operand sg_layout = [8, 4] sg_data = [32, 64] : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// CHECK-LABEL: @set_op_layout_attr_not_anchor_op
+func.func @set_op_layout_attr_not_anchor_op(%arg0: memref<4096x4096xf16>) {
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>
+ %2 = arith.extf %1 : vector<256x32xf16> to vector<256x32xf32> // expected-note {{target op}}
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["arith.extf"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ // expected-error at below {{Cannot set anchor layout to op: arith.extf}}
transform.xegpu.set_op_layout_attr %0 sg_layout = [8, 4] sg_data = [32, 64] : !transform.any_op
transform.yield
}
diff --git a/mlir/test/Dialect/XeGPU/transform-ops.mlir b/mlir/test/Dialect/XeGPU/transform-ops.mlir
index 13ed24ebf0a3a..fce0012878eeb 100644
--- a/mlir/test/Dialect/XeGPU/transform-ops.mlir
+++ b/mlir/test/Dialect/XeGPU/transform-ops.mlir
@@ -264,7 +264,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%0 = transform.structured.match ops{["arith.extf"]} in %arg1 : (!transform.any_op) -> !transform.any_op
// CHECK: transform.xegpu.set_op_layout_attr %{{.*}}
- transform.xegpu.set_op_layout_attr %0 sg_layout = [8, 4] sg_data = [32, 64] : !transform.any_op
+ transform.xegpu.set_op_layout_attr %0 operand sg_layout = [8, 4] sg_data = [32, 64] : !transform.any_op
transform.yield
}
}
@@ -287,7 +287,27 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%0 = transform.structured.match ops{["arith.addf"]} in %arg1 : (!transform.any_op) -> !transform.any_op
// CHECK: transform.xegpu.set_op_layout_attr %{{.*}}
- transform.xegpu.set_op_layout_attr %0 index = 1 sg_layout = [8, 4] sg_data = [32, 64] inst_data = [8, 16] : !transform.any_op
+ transform.xegpu.set_op_layout_attr %0 operand index = 1 sg_layout = [8, 4] sg_data = [32, 64] inst_data = [8, 16] : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// CHECK-LABEL: @set_op_layout_attr_anchor
+func.func @set_op_layout_attr_anchor(%arg0: memref<4096x4096xf16>) {
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>
+ // CHECK: = xegpu.load_nd %0[0, 0]
+ // CHECK-SAME: <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], inst_data = [8, 16]>}>
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["xegpu.load_nd"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ // CHECK: transform.xegpu.set_op_layout_attr %{{.*}}
+ transform.xegpu.set_op_layout_attr %0 index = 0 sg_layout = [8, 4] sg_data = [32, 64] inst_data = [8, 16] : !transform.any_op
transform.yield
}
}
diff --git a/mlir/test/python/dialects/transform_xegpu_ext.py b/mlir/test/python/dialects/transform_xegpu_ext.py
index 2b11acb04ed5b..d8d2b9cb33ca0 100644
--- a/mlir/test/python/dialects/transform_xegpu_ext.py
+++ b/mlir/test/python/dialects/transform_xegpu_ext.py
@@ -97,10 +97,12 @@ def setOpLayoutAttrOperandMinimal():
sequence.bodyTarget,
sg_layout=[6, 4],
sg_data=[32, 16],
+ operand=True,
)
transform.YieldOp()
# CHECK-LABEL: TEST: setOpLayoutAttr
# CHECK: transform.xegpu.set_op_layout_attr %
+ # CHECK: operand
# NO-CHECK: index = 0
# NO-CHECK: result
# CHECK: sg_layout = [6, 4]
@@ -127,8 +129,8 @@ def setOpLayoutAttrResult():
transform.YieldOp()
# CHECK-LABEL: TEST: setOpLayoutAttrResult
# CHECK: transform.xegpu.set_op_layout_attr %
- # NO-CHECK: index = 0
# CHECK: result
+ # NO-CHECK: index = 0
# CHECK: sg_layout = [6, 4]
# CHECK: sg_data = [32, 16]
# CHECK: inst_data = [8, 16]
@@ -154,14 +156,40 @@ def setOpLayoutAttrResultSlice():
transform.YieldOp()
# CHECK-LABEL: TEST: setOpLayoutAttrResultSlice
# CHECK: transform.xegpu.set_op_layout_attr %
- # NO-CHECK: index = 0
# CHECK: result
+ # NO-CHECK: index = 0
# CHECK: sg_layout = [6, 4]
# CHECK: sg_data = [32, 16]
# CHECK: inst_data = [8, 16]
# CHECK: slice_dims = [0]
+ at run
+def setOpLayoutAttrAnchor():
+ sequence = transform.SequenceOp(
+ transform.FailurePropagationMode.Propagate,
+ [],
+ transform.OperationType.get("xegpu.dpas"),
+ )
+ with InsertionPoint(sequence.body):
+ xegpu.set_op_layout_attr(
+ sequence.bodyTarget,
+ index=0,
+ sg_layout=[6, 4],
+ sg_data=[32, 16],
+ inst_data=[8, 16],
+ )
+ transform.YieldOp()
+ # CHECK-LABEL: TEST: setOpLayoutAttrAnchor
+ # CHECK: transform.xegpu.set_op_layout_attr %
+ # NO-CHECK: result
+ # NO-CHECK: operand
+ # NO-CHECK: index = 0
+ # CHECK: sg_layout = [6, 4]
+ # CHECK: sg_data = [32, 16]
+ # CHECK: inst_data = [8, 16]
+
+
@run
def setGPULaunchThreadsOp():
sequence = transform.SequenceOp(
>From 0d981c41fa3826ef26628a9e1801d3c9fa21a3da Mon Sep 17 00:00:00 2001
From: Tuomas Karna <tuomas.karna at intel.com>
Date: Wed, 17 Dec 2025 10:21:37 +0200
Subject: [PATCH 2/4] add verifier and fix python tests
---
.../XeGPU/TransformOps/XeGPUTransformOps.td | 2 ++
.../XeGPU/TransformOps/XeGPUTransformOps.cpp | 12 +++++++-----
mlir/test/python/dialects/transform_xegpu_ext.py | 16 ++++++++--------
3 files changed, 17 insertions(+), 13 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td b/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
index 8c5539ad19382..23dabe4eb380a 100644
--- a/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
@@ -173,6 +173,8 @@ def SetOpLayoutAttrOp : Op<Transform_Dialect, "xegpu.set_op_layout_attr", [
return getMixedValues(getStaticInstData(), getInstData(), b);
}
}];
+
+ let hasVerifier = 1;
}
def SetGPULaunchThreadsOp
diff --git a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
index e4bcff4b996eb..f785b91a37d37 100644
--- a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
@@ -325,11 +325,6 @@ transform::SetOpLayoutAttrOp::apply(transform::TransformRewriter &rewriter,
bool resultTarget = getResult();
bool operandTarget = getOperand();
- if (resultTarget && operandTarget) {
- return emitSilenceableFailure(getLoc())
- << "Cannot set both result and operand layout attributes.";
- }
-
int64_t index = getIndex();
if (resultTarget && index >= target->getNumResults()) {
return emitSilenceableFailure(getLoc())
@@ -388,6 +383,13 @@ void transform::SetOpLayoutAttrOp::getEffects(
modifiesPayload(effects);
}
+LogicalResult transform::SetOpLayoutAttrOp::verify() {
+ if (getResult() && getOperand()) {
+ return emitOpError("Cannot set both result and operand simultaneously.");
+ }
+ return success();
+}
+
void transform::SetGPULaunchThreadsOp::build(
OpBuilder &builder, OperationState &ostate, Value target,
ArrayRef<OpFoldResult> mixedThreads) {
diff --git a/mlir/test/python/dialects/transform_xegpu_ext.py b/mlir/test/python/dialects/transform_xegpu_ext.py
index d8d2b9cb33ca0..e3e1313cf5f81 100644
--- a/mlir/test/python/dialects/transform_xegpu_ext.py
+++ b/mlir/test/python/dialects/transform_xegpu_ext.py
@@ -103,11 +103,11 @@ def setOpLayoutAttrOperandMinimal():
# CHECK-LABEL: TEST: setOpLayoutAttr
# CHECK: transform.xegpu.set_op_layout_attr %
# CHECK: operand
- # NO-CHECK: index = 0
- # NO-CHECK: result
+ # CHECK-NOT: index = 0
+ # CHECK-NOT: result
# CHECK: sg_layout = [6, 4]
# CHECK: sg_data = [32, 16]
- # NO-CHECK: inst_data
+ # CHECK-NOT: inst_data
@run
@@ -130,7 +130,7 @@ def setOpLayoutAttrResult():
# CHECK-LABEL: TEST: setOpLayoutAttrResult
# CHECK: transform.xegpu.set_op_layout_attr %
# CHECK: result
- # NO-CHECK: index = 0
+ # CHECK-NOT: index = 0
# CHECK: sg_layout = [6, 4]
# CHECK: sg_data = [32, 16]
# CHECK: inst_data = [8, 16]
@@ -157,7 +157,7 @@ def setOpLayoutAttrResultSlice():
# CHECK-LABEL: TEST: setOpLayoutAttrResultSlice
# CHECK: transform.xegpu.set_op_layout_attr %
# CHECK: result
- # NO-CHECK: index = 0
+ # CHECK-NOT: index = 0
# CHECK: sg_layout = [6, 4]
# CHECK: sg_data = [32, 16]
# CHECK: inst_data = [8, 16]
@@ -182,9 +182,9 @@ def setOpLayoutAttrAnchor():
transform.YieldOp()
# CHECK-LABEL: TEST: setOpLayoutAttrAnchor
# CHECK: transform.xegpu.set_op_layout_attr %
- # NO-CHECK: result
- # NO-CHECK: operand
- # NO-CHECK: index = 0
+ # CHECK-NOT: result
+ # CHECK-NOT: operand
+ # CHECK-NOT: index = 0
# CHECK: sg_layout = [6, 4]
# CHECK: sg_data = [32, 16]
# CHECK: inst_data = [8, 16]
>From 99dda99c93c2bcfe882a07f8203367153c0c391f Mon Sep 17 00:00:00 2001
From: Tuomas Karna <tuomas.karna at intel.com>
Date: Thu, 18 Dec 2025 15:50:29 +0200
Subject: [PATCH 3/4] use AnchorLayoutInterface
---
.../XeGPU/TransformOps/XeGPUTransformOps.cpp | 17 +++++++----------
1 file changed, 7 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
index f785b91a37d37..5ea55755f22d7 100644
--- a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
@@ -350,26 +350,23 @@ transform::SetOpLayoutAttrOp::apply(transform::TransformRewriter &rewriter,
getContext(), layout, DenseI64ArrayAttr::get(getContext(), sliceDims));
}
- // Set layout attribute for the op result or operand
+ // Set layout attribute
if (resultTarget) {
+ // for the op result
xegpu::setDistributeLayoutAttr(target->getResult(index), layout);
} else if (operandTarget) {
+ // or operand
xegpu::setDistributeLayoutAttr(target->getOpOperand(index), layout);
} else {
- // Set anchor layout if requested.
- // TODO use AnchorLayoutInterface when available.
- if (!isa<xegpu::LoadNdOp>(target)) {
+ // or anchor layout.
+ auto anchorOp = dyn_cast<xegpu::AnchorLayoutInterface>(target);
+ if (!anchorOp) {
auto diag = emitSilenceableFailure(getLoc())
<< "Cannot set anchor layout to op: " << target->getName();
diag.attachNote(target->getLoc()) << "target op";
return diag;
}
- auto loadOp = dyn_cast<xegpu::LoadNdOp>(target);
- if (loadOp)
- loadOp.setLayoutAttr(layout);
- auto prefetchOp = dyn_cast<xegpu::PrefetchOp>(target);
- if (prefetchOp)
- prefetchOp.setLayoutAttr(layout);
+ anchorOp.setAnchorLayout(layout);
}
return DiagnosedSilenceableFailure::success();
}
>From 129443c3c1d9eed43ca0a74ecce7f4941ec673de Mon Sep 17 00:00:00 2001
From: Tuomas Karna <tuomas.karna at intel.com>
Date: Tue, 30 Dec 2025 17:39:31 +0200
Subject: [PATCH 4/4] dpas op requires anchor layout for A,B, and C tiles
---
.../XeGPU/TransformOps/XeGPUTransformOps.cpp | 20 ++++++++++++++++---
1 file changed, 17 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
index 5ea55755f22d7..7bc67da8263dc 100644
--- a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
@@ -352,13 +352,27 @@ transform::SetOpLayoutAttrOp::apply(transform::TransformRewriter &rewriter,
// Set layout attribute
if (resultTarget) {
- // for the op result
+ // op result
xegpu::setDistributeLayoutAttr(target->getResult(index), layout);
} else if (operandTarget) {
- // or operand
+ // op operand
xegpu::setDistributeLayoutAttr(target->getOpOperand(index), layout);
+ } else if (auto dpasOp = dyn_cast<xegpu::DpasOp>(target)) {
+ // dpas op is a special case where layout needs to be set for A, B, and C
+ if (index == 0)
+ dpasOp.getProperties().layout_a = layout;
+ else if (index == 1)
+ dpasOp.getProperties().layout_b = layout;
+ else if (index == 2)
+ dpasOp.getProperties().layout_cd = layout;
+ else {
+ auto diag = emitSilenceableFailure(getLoc())
+ << "Invalid index for setting dpas op layout: " << index;
+ diag.attachNote(target->getLoc()) << "target op";
+ return diag;
+ }
} else {
- // or anchor layout.
+ // op's anchor layout.
auto anchorOp = dyn_cast<xegpu::AnchorLayoutInterface>(target);
if (!anchorOp) {
auto diag = emitSilenceableFailure(getLoc())
More information about the Mlir-commits
mailing list