[Mlir-commits] [mlir] [mlir][xegpu] Handle size-1 broadcast dim in insert_strided_slice layout setup (PR #211725)

[Jianhui Li]

https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/211725

This PR enhance `setupInsertStridedSliceResultLayout` to treat a size-1 source dim as a broadcast dim: keep its lane_data at 1 instead of requiring divisibility. 

Add regression tests for both the layout propagation setup and the subgroup-to-lane distribution of such an op.

Assisted-by-claude

>From 59b316d9c9fbd24e0f0ffb75aa3e9e41e11d416d Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 24 Jul 2026 04:34:32 +0000
Subject: [PATCH] [mlir][xegpu] Handle size-1 broadcast dim in
 insert_strided_slice layout setup

setupInsertStridedSliceResultLayout asserted that every source dimension of
a vector.insert_strided_slice divides the consumer lane layout. This fails
for a size-1 source dim distributed across multiple lanes (e.g. inserting
vector<1x1xf32> into vector<8x1xf32> with lane_layout = [1, 16]), even though
that dim is simply broadcast across the lanes and the slice is inserted along
another dim.

Treat a size-1 source dim as a broadcast dim: keep its lane_data at 1 instead
of requiring divisibility. Add regression tests for both the layout
propagation setup and the subgroup-to-lane distribution of such an op.

Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
 .../XeGPU/Transforms/XeGPULayoutImpl.cpp      | 13 ++++++++----
 mlir/test/Dialect/XeGPU/propagate-layout.mlir | 17 +++++++++++++++
 .../XeGPU/sg-to-lane-distribute-unit.mlir     | 21 +++++++++++++++++++
 3 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 48889a8b8febc..ea0846d9d9da4 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -2683,10 +2683,15 @@ xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
                     "insertStridedSlice.");
   } else if (layoutKind == xegpu::LayoutKind::Lane) {
     for (int dim = 0; dim < srcRank; dim++) {
-      assert(srcShape[dim] % consumerLaneLayout[dim] == 0 &&
-             "srcShape must be divisible by laneLayout for all dimensions");
-      laneDataValue = std::min(srcShape[dim] / consumerLaneLayout[dim],
-                               consumerLaneData[dim]);
+      // A size-1 source dim is broadcast across the lanes of that dim.
+      if (srcShape[dim] == 1) {
+        laneDataValue = 1;
+      } else {
+        assert(srcShape[dim] % consumerLaneLayout[dim] == 0 &&
+               "srcShape must be divisible by laneLayout for all dimensions");
+        laneDataValue = std::min(srcShape[dim] / consumerLaneLayout[dim],
+                                 consumerLaneData[dim]);
+      }
       requiredResLayout =
           requiredResLayout.setDimData(dim, -1, -1, laneDataValue);
     }
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index 79a5d229263c5..d2992904ab495 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -960,6 +960,23 @@ func.func @insert_strided_slice_with_slice_layout(%arg0: memref<8x16xf32>) {
 }
 }
 
+// -----
+gpu.module @test {
+// CHECK-LABEL: func.func @insert_strided_slice_lane_broadcast_dim(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x1xf32>) {
+// CHECK: %[[CST_SMALL:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} dense<1.000000e+00> : vector<1x1xf32>
+// CHECK: %[[CST_LARGE:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<8x1xf32>
+// CHECK: %[[INSERT:.*]] = vector.insert_strided_slice %[[CST_SMALL]], %[[CST_LARGE]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, offsets = [0, 0], strides = [1, 1]} : vector<1x1xf32> into vector<8x1xf32>
+func.func @insert_strided_slice_lane_broadcast_dim(%arg0: memref<8x1xf32>) {
+  %cst_small = arith.constant dense<1.0> : vector<1x1xf32>
+  %cst_large = arith.constant dense<0.0> : vector<8x1xf32>
+  %insert = vector.insert_strided_slice %cst_small, %cst_large {offsets = [0, 0], strides = [1, 1]} : vector<1x1xf32> into vector<8x1xf32>
+  %tdesc = xegpu.create_nd_tdesc %arg0 : memref<8x1xf32> -> !xegpu.tensor_desc<8x1xf32>
+  xegpu.store_nd %insert, %tdesc[0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x1xf32>, !xegpu.tensor_desc<8x1xf32>
+  return
+}
+}
+
 // -----
 gpu.module @test{
   // CHECK-LABEL: load_store_matrix
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
index fe356e6af35c1..bdd45a5d5fcea 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -1024,6 +1024,27 @@ gpu.func @vector_insert_strided_slice_different_ranks() {
   gpu.return
 }
 
+// The lane-distributed dim (dim 1) is size 1 in the source and thus broadcast
+// across the lanes of that dim; the insert happens on dim 0 only, so the
+// distributed shapes are unchanged.
+// CHECK-LABEL: gpu.func @vector_insert_strided_slice_lane_broadcast_dim
+// CHECK: %[[ISS:.*]] = vector.insert_strided_slice %{{.*}}, %{{.*}} {offsets = [0, 0], strides = [1, 1]} : vector<1x1xf32> into vector<8x1xf32>
+gpu.func @vector_insert_strided_slice_lane_broadcast_dim() {
+  %0 = "some_op"()
+    : () -> vector<1x1xf32>
+  %1 = "some_op"()
+    : () -> vector<8x1xf32>
+  %2 = vector.insert_strided_slice %0, %1 { offsets = [0, 0], strides = [1, 1]
+    }
+    : vector<1x1xf32> into vector<8x1xf32>
+  %cl2 = xegpu.convert_layout %2
+    <{
+      input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+      target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+    }> : vector<8x1xf32>
+  gpu.return
+}
+
 // CHECK-LABEL: gpu.func @convert_layout_removed_when_compatible
 // CHECK-NOT: xegpu.convert_layout
 gpu.func @convert_layout_removed_when_compatible() {