[Mlir-commits] [mlir] [MLIR][XeGPU] Make shape cast respect full replication layout semantics (PR #229395)

[Artem Kroviakov]

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

>From 4e59165461852aadbf7ba5265ca2927839659275 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Wed, 7 Oct 2026 10:56:22 +0000
Subject: [PATCH] [MLIR][XeGPU] Make shape cast respect full replication layout
 semantics

---
 .../include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h |  8 ++------
 mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp       |  7 ++++++-
 .../Dialect/XeGPU/propagate-layout-inst-data.mlir | 15 +++++++++++++++
 3 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index 4db69bda113dc..ab526793f2aa8 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -17,7 +17,6 @@
 #include <optional>
 
 namespace mlir {
-
 class UnrealizedConversionCastOp;
 class VectorType;
 class OpOperand;
@@ -38,7 +37,6 @@ struct uArch;
 } // namespace xegpu
 
 namespace xegpu {
-
 /// Flatten a set of ValueRange into a single SmallVector<Value>
 SmallVector<Value> flattenValues(ArrayRef<ValueRange> values);
 
@@ -247,8 +245,8 @@ bool hasStaticShapeAndStrides(MemRefType type);
 bool matchUnitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
                            SmallVector<int64_t> &expandedUnitDims);
 
-// Checks if dst shape is an expansion of src shape where each dimension in src
-// is split into one or more consecutive dimensions in dst
+// Checks if dst shape is a non-unit expansion of src shape where each
+// dimension in src is split into one or more consecutive dimensions in dst.
 bool matchSplitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
                             SmallVector<SmallVector<int64_t>> &splitDimGroups);
 
@@ -304,9 +302,7 @@ void cleanupUnrealizedConversionCasts(
 // dst=[1,4096] -> true, collapseDims=[[],[0,1,2]].
 bool matchDimCollapse(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
                       SmallVector<SmallVector<int64_t>> &collapseDims);
-
 } // namespace xegpu
-
 } // namespace mlir
 
 #endif // MLIR_DIALECT_XEGPU_UTILS_XEGPUUTILS_H_
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index d1d4f044c8ed0..e22364df9e590 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -864,10 +864,15 @@ bool xegpu::matchUnitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
 // is split into one or more consecutive dimensions in dst whose product equals
 // the original dimension. Populates splitDimGroups with groups of dst indices
 // that correspond to each src dimension. Example: src=[6,4], dst=[2,3,2,2] ->
-// true
+// true. Shapes that only insert unit dimensions are  rejected here.
 bool xegpu::matchSplitDimExpansion(
     ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
     SmallVector<SmallVector<int64_t>> &splitDimGroups) {
+  SmallVector<int64_t> expandedUnitDims;
+  if (matchUnitDimExpansion(src, dst, expandedUnitDims) &&
+      !expandedUnitDims.empty())
+    return false;
+
   // each dim in src can be mapped to one or more dims in dst whose product
   // equals to the src dim
   size_t srcIdx = 0;
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 30e02cc42095c..a0b261f12402d 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -772,3 +772,18 @@ func.func @shape_cast_split_fills_strided_inner_dim(%arg0: memref<16x1024xbf16>)
   return
 }
 }
+
+// -----
+gpu.module @test {
+// CHECK-LABEL: func.func @shape_cast_insert_unit_dim_replicated_lanes(
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 4], lane_layout = [1, 4, 4], lane_data = [1, 1, 1]>, dims = [1]>} dense<0.000000e+00> : vector<8x32xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[CST]] {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 4], lane_layout = [1, 4, 4], lane_data = [1, 1, 1]>} : vector<8x32xf16> to vector<8x1x32xf16>
+func.func @shape_cast_insert_unit_dim_replicated_lanes() {
+  %cst = arith.constant dense<0.000000e+00> : vector<8x32xf16>
+  %0 = vector.shape_cast %cst : vector<8x32xf16> to vector<8x1x32xf16>
+  %1 = xegpu.convert_layout %0
+     <{target_layout = #xegpu.layout<inst_data = [1, 1, 4], lane_layout = [1, 4, 4], lane_data = [1, 1, 1]>}>
+     : vector<8x1x32xf16>
+  return
+}
+}