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

[llvmlistbot at llvm.org]

Author: Jianhui Li
Date: 2026-07-30T09:20:53-07:00
New Revision: 665d200b160d15ef7e8ddeb21c7028d774c7e1b1

URL: https://github.com/llvm/llvm-project/commit/665d200b160d15ef7e8ddeb21c7028d774c7e1b1
DIFF: https://github.com/llvm/llvm-project/commit/665d200b160d15ef7e8ddeb21c7028d774c7e1b1.diff

LOG: [mlir][xegpu] Handle size-1 broadcast dim in insert_strided_slice layout setup (#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

Co-authored-by: Claude Opus 4.8 <noreply at anthropic.com>

Added: 
    

Modified: 
    mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
    mlir/test/Dialect/XeGPU/propagate-layout.mlir
    mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir

Removed: 
    


################################################################################
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 d48cb7993bfab..5d45ec79a5a8a 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -959,6 +959,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 a338ed5de7efc..a7fe2d9534351 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -987,6 +987,27 @@ gpu.func @vector_insert_strided_slice_
diff erent_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() {