[Mlir-commits] [mlir] [MLIR][XeGPU] Get correct distribute layout from multi-result ops. (PR #198028)
Sang Ik Lee
llvmlistbot at llvm.org
Thu May 21 09:59:09 PDT 2026
https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/198028
>From 2e8e124d93ece9d2a953420fa9662bce37c24154 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Fri, 15 May 2026 21:20:19 +0000
Subject: [PATCH 1/3] [MLIR][XeGPU] Get correct distribute layout from
multi-result ops. Ops like vector.deinterleave produces more than one
results. Pass correct result to xegpu::getDistributeLayoutAttr() for such
ops.
---
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp | 9 ++++++++-
1 file changed, 8 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 813e9608dbf8e..92d5f7de02d90 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1910,7 +1910,14 @@ xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
// For non-anchor ops, derive the operand layout from the op's result
// layout via op-specific semantics.
xegpu::DistributeLayoutAttr resLayout;
- if (op->getNumResults() == 1)
+ if (op->getNumResults() == 1) {
resLayout = xegpu::getDistributeLayoutAttr(op->getResult(0));
+ } else {
+ Value val = operand.get();
+ if (auto opResult = dyn_cast<OpResult>(val)) {
+ auto resultIdx = opResult.getResultNumber();
+ resLayout = xegpu::getDistributeLayoutAttr(op->getResult(resultIdx));
+ }
+ }
return inferSourceLayoutFromResultForNonAnchorOp(operand, resLayout);
}
>From 7268c706a6499b0098c1b1bc8d4823e6c4e33daf Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Tue, 19 May 2026 18:04:06 +0000
Subject: [PATCH 2/3] Add test cases.
---
.../XeGPU/resolve-layout-conflicts.mlir | 59 +++++++++++++++++++
1 file changed, 59 insertions(+)
diff --git a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
index 449afb7b8378b..123c13b0ebe55 100644
--- a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
+++ b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
@@ -292,6 +292,65 @@ func.func @extract_source_conflict_with_order() -> vector<16x32xf16> {
%1 = vector.extract %0[0, 0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [0, 1]>} : vector<16x32xf16> from vector<2x4x16x32xf16>
return %1 : vector<16x32xf16>
}
+
+// vector.deinterleave is the producer with 2 results. A second
+// vector.deinterleave (also a non-anchor op with 2 results) consumes result #1
+// of the first. getConsumerLayoutAt must index into the consumer op's results
+// using the producer result number (resultIdx=1) rather than always result(0).
+//
+// No conflict: result #1 of the first deinterleave (inst_data=[8,16]) matches
+// inferDeinterleaveSourceLayout(inst_data=[8,8]) = inst_data=[8,16].
+// CHECK-LABEL: func.func @nested_deinterleave_no_conflict
+// CHECK: %[[X:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 32]>} : () -> vector<8x32xf16>
+// CHECK-NEXT: %[[A0:.*]], %[[A1:.*]] = vector.deinterleave %[[X]]
+// CHECK-SAME: {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>, layout_result_1 = #xegpu.layout<inst_data = [8, 16]>}
+// CHECK-SAME: : vector<8x32xf16> -> vector<8x16xf16>
+// CHECK-NEXT: %{{.*}}, %{{.*}} = vector.deinterleave %[[A1]]
+// CHECK-SAME: {layout_result_0 = #xegpu.layout<inst_data = [8, 8]>, layout_result_1 = #xegpu.layout<inst_data = [8, 8]>}
+// CHECK-SAME: : vector<8x16xf16> -> vector<8x8xf16>
+// CHECK-NOT: xegpu.convert_layout
+// CHECK: return
+func.func @nested_deinterleave_no_conflict() {
+ %x = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 32]>} : () -> vector<8x32xf16>
+ %a0, %a1 = vector.deinterleave %x {
+ layout_result_0 = #xegpu.layout<inst_data = [8, 16]>,
+ layout_result_1 = #xegpu.layout<inst_data = [8, 16]>
+ } : vector<8x32xf16> -> vector<8x16xf16>
+ %b0, %b1 = vector.deinterleave %a1 {
+ layout_result_0 = #xegpu.layout<inst_data = [8, 8]>,
+ layout_result_1 = #xegpu.layout<inst_data = [8, 8]>
+ } : vector<8x16xf16> -> vector<8x8xf16>
+ return
+}
+
+// Conflict: result #1 of the first deinterleave carries inst_data=[16,16] but
+// the second deinterleave expects inst_data=[8,16] as its source (inferred from
+// inst_data=[8,8] via inferDeinterleaveSourceLayout). A convert_layout must be
+// inserted to bridge the mismatch.
+// CHECK-LABEL: func.func @nested_deinterleave_with_conflict
+// CHECK: %[[X:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 32]>} : () -> vector<16x32xf16>
+// CHECK-NEXT: %[[A0:.*]], %[[A1:.*]] = vector.deinterleave %[[X]]
+// CHECK-SAME: {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>, layout_result_1 = #xegpu.layout<inst_data = [16, 16]>}
+// CHECK-SAME: : vector<16x32xf16> -> vector<16x16xf16>
+// CHECK-NEXT: %[[CVT:.*]] = xegpu.convert_layout %[[A1]]
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [16, 16]>, target_layout = #xegpu.layout<inst_data = [8, 16]>}>
+// CHECK-SAME: : vector<16x16xf16>
+// CHECK-NEXT: %{{.*}}, %{{.*}} = vector.deinterleave %[[CVT]]
+// CHECK-SAME: {layout_result_0 = #xegpu.layout<inst_data = [8, 8]>, layout_result_1 = #xegpu.layout<inst_data = [8, 8]>}
+// CHECK-SAME: : vector<16x16xf16> -> vector<16x8xf16>
+// CHECK: return
+func.func @nested_deinterleave_with_conflict() {
+ %x = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 32]>} : () -> vector<16x32xf16>
+ %a0, %a1 = vector.deinterleave %x {
+ layout_result_0 = #xegpu.layout<inst_data = [8, 16]>,
+ layout_result_1 = #xegpu.layout<inst_data = [16, 16]>
+ } : vector<16x32xf16> -> vector<16x16xf16>
+ %b0, %b1 = vector.deinterleave %a1 {
+ layout_result_0 = #xegpu.layout<inst_data = [8, 8]>,
+ layout_result_1 = #xegpu.layout<inst_data = [8, 8]>
+ } : vector<16x16xf16> -> vector<16x8xf16>
+ return
+}
}
// -----
>From e0b0c71c4749459ccce6c88192c822c6513446e4 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 20 May 2026 17:03:57 +0000
Subject: [PATCH 3/3] New logic in else clause is general. Remove if clause.
---
.../lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp | 12 ++++--------
1 file changed, 4 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 6df8f5b701593..322cfb7b0002c 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1920,14 +1920,10 @@ xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
// For non-anchor ops, derive the operand layout from the op's result
// layout via op-specific semantics.
xegpu::DistributeLayoutAttr resLayout;
- if (op->getNumResults() == 1) {
- resLayout = xegpu::getDistributeLayoutAttr(op->getResult(0));
- } else {
- Value val = operand.get();
- if (auto opResult = dyn_cast<OpResult>(val)) {
- auto resultIdx = opResult.getResultNumber();
- resLayout = xegpu::getDistributeLayoutAttr(op->getResult(resultIdx));
- }
+ Value val = operand.get();
+ if (auto opResult = dyn_cast<OpResult>(val)) {
+ auto resultIdx = opResult.getResultNumber();
+ resLayout = xegpu::getDistributeLayoutAttr(op->getResult(resultIdx));
}
return inferSourceLayoutFromResultForNonAnchorOp(operand, resLayout);
}
More information about the Mlir-commits
mailing list