[Mlir-commits] [mlir] [MLIR][XeGPU] Fix insert_strided_slice op in subgroup distribution (PR #180604)

[Jianhui Li]

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

The PR modifies the subgroup distribution pass to only sink insert_strided_slice operation if it becomes the last op before yield.  It avoids sinking insert_strided_slice multiple times and cause potential issue in worst case. 

>From 82bccdcedea54a46ebc61a89005867d08aa32979 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Mon, 9 Feb 2026 19:50:57 +0000
Subject: [PATCH] fix insert_strided_slice in sg distribution

---
 .../Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp   | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
index 5cd9772204590..aa1dfaa9e0fda 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
@@ -1757,8 +1757,11 @@ struct VectorInsertStridedSliceDistribution
   using gpu::WarpDistributionPattern::WarpDistributionPattern;
   LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,
                                 PatternRewriter &rewriter) const override {
-    OpOperand *operand =
-        getWarpResult(warpOp, llvm::IsaPred<vector::InsertStridedSliceOp>);
+    OpOperand *operand = getWarpResult(warpOp, [&](Operation *op) {
+      // Check if the InsertStridedSliceOp is the last op before yield op
+      return llvm::IsaPred<vector::InsertStridedSliceOp>(op) &&
+             warpOp.getTerminator()->getPrevNode() == op;
+    });
     if (!operand)
       return failure();
     unsigned int operandNumber = operand->getOperandNumber();