[Mlir-commits] [mlir] [MLIR][XeGPU] Allow uniform vectors in layout conflict resolution (PR #183756)

[Artem Kroviakov]

https://github.com/akroviakov updated https://github.com/llvm/llvm-project/pull/183756

>From b805dc4eac08f48c9686cc8a0372003fdbbfb0e1 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Fri, 27 Feb 2026 15:40:58 +0000
Subject: [PATCH 1/2] [MLIR][XeGPU] Allow uniform vectors in layout conflict
 resolution

---
 .../XeGPU/Transforms/XeGPUPropagateLayout.cpp        |  3 +--
 .../test/Dialect/XeGPU/resolve-layout-conflicts.mlir | 12 ++++++++++++
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 10cd65b080405..99abb3d22f9cd 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1311,8 +1311,7 @@ ResolveLayoutConflicts::resolveVectorConsumer(OpOperand &operand) {
   // Get the current layout of the vector value.
   auto producerLayout = xegpu::getDistributeLayoutAttr(vectorValue);
   if (!producerLayout)
-    return consumerOp->emitError("Vector operand has no layout assigned.");
-
+    return success(); // vector with no layout is uniform
   // Get the consumer expected layout at this operand.
   auto consumerLayout = xegpu::getConsumerLayoutAt(operand);
   if (!consumerLayout)
diff --git a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
index 7f9796ae2aecf..8a7d85347f920 100644
--- a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
+++ b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
@@ -100,6 +100,18 @@ func.func @elementwise_conflict() -> vector<32x32xf16> {
   return %2 : vector<32x32xf16>
 }
 
+// CHECK-LABEL: func.func @elementwise_conflict_uniform
+// CHECK-DAG:     %[[V0:.*]] = "some_op"() : () -> vector<32x32xf16>
+// CHECK-DAG:     %[[V1:.*]] = "some_op"() : () -> vector<32x32xf16>
+// CHECK:         %[[ADD:.*]] = arith.addf %[[V0]], %[[V1]] : vector<32x32xf16>
+// CHECK:         return %[[ADD]] : vector<32x32xf16>
+func.func @elementwise_conflict_uniform() -> vector<32x32xf16> {
+  %0 = "some_op"() : () -> vector<32x32xf16>
+  %1 = "some_op"() : () -> vector<32x32xf16>
+  %2 = arith.addf %0, %1 : vector<32x32xf16>
+  return %2 : vector<32x32xf16>
+}
+
 // CHECK-LABEL: func.func @broadcast_source_conflict
 // CHECK:         %[[V0:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [16]>} : () -> vector<16xf16>
 // CHECK:         %[[CVT:.*]] = xegpu.convert_layout %[[V0]]

>From abbe33a5b598d52b0b3e94c58850dbb89ea9b55e Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Tue, 3 Mar 2026 09:56:42 +0000
Subject: [PATCH 2/2] Add warning for non-1D vectors

---
 .../XeGPU/Transforms/XeGPUPropagateLayout.cpp  |  8 ++++++--
 .../XeGPU/resolve-layout-conflicts.mlir        | 18 +++++++++---------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 99abb3d22f9cd..87835fb191604 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1310,8 +1310,12 @@ ResolveLayoutConflicts::resolveVectorConsumer(OpOperand &operand) {
   Operation *consumerOp = operand.getOwner();
   // Get the current layout of the vector value.
   auto producerLayout = xegpu::getDistributeLayoutAttr(vectorValue);
-  if (!producerLayout)
-    return success(); // vector with no layout is uniform
+  if (!producerLayout) {
+    if (auto vectorTy = dyn_cast<VectorType>(vectorValue.getType());
+        vectorTy && vectorTy.getRank() > 1)
+      consumerOp->emitWarning("Expected layout for non-1D vectors.");
+    return success(); // uniform non-tensor-data vector does not require layout
+  }
   // Get the consumer expected layout at this operand.
   auto consumerLayout = xegpu::getConsumerLayoutAt(operand);
   if (!consumerLayout)
diff --git a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
index 8a7d85347f920..c73a9990542de 100644
--- a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
+++ b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
@@ -101,15 +101,15 @@ func.func @elementwise_conflict() -> vector<32x32xf16> {
 }
 
 // CHECK-LABEL: func.func @elementwise_conflict_uniform
-// CHECK-DAG:     %[[V0:.*]] = "some_op"() : () -> vector<32x32xf16>
-// CHECK-DAG:     %[[V1:.*]] = "some_op"() : () -> vector<32x32xf16>
-// CHECK:         %[[ADD:.*]] = arith.addf %[[V0]], %[[V1]] : vector<32x32xf16>
-// CHECK:         return %[[ADD]] : vector<32x32xf16>
-func.func @elementwise_conflict_uniform() -> vector<32x32xf16> {
-  %0 = "some_op"() : () -> vector<32x32xf16>
-  %1 = "some_op"() : () -> vector<32x32xf16>
-  %2 = arith.addf %0, %1 : vector<32x32xf16>
-  return %2 : vector<32x32xf16>
+// CHECK-DAG:     %[[V0:.*]] = "some_op"() : () -> vector<2xf16>
+// CHECK-DAG:     %[[V1:.*]] = "some_op"() : () -> vector<2xf16>
+// CHECK:         %[[ADD:.*]] = arith.addf %[[V0]], %[[V1]] : vector<2xf16>
+// CHECK:         return %[[ADD]] : vector<2xf16>
+func.func @elementwise_conflict_uniform() -> vector<2xf16> {
+  %0 = "some_op"() : () -> vector<2xf16>
+  %1 = "some_op"() : () -> vector<2xf16>
+  %non_tensor_data_vec = arith.addf %0, %1 : vector<2xf16>
+  return %non_tensor_data_vec : vector<2xf16>
 }
 
 // CHECK-LABEL: func.func @broadcast_source_conflict