[Mlir-commits] [mlir] [mlir][xegpu] Fill split-group inner dims on shape_cast result layout (PR #227560)

[Jianhui Li]

================
@@ -2752,6 +2752,44 @@ xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
                                            resShape[innerMostDim], uArch);
 }
 
+xegpu::DistributeLayoutAttr xegpu::setupShapeCastResultLayout(
+    xegpu::LayoutKind layoutKind, VectorType srcVecTy, VectorType resVecTy,
+    DistributeLayoutAttr consumerLayout) {
+  // TODO: work out the subgroup-level rule; leave such layouts alone for now.
+  if (!consumerLayout || layoutKind == xegpu::LayoutKind::Subgroup)
+    return consumerLayout;
+
+  ArrayRef<int64_t> resShape = resVecTy.getShape();
+  SmallVector<SmallVector<int64_t>> splitDimGroups;
+  if (!xegpu::matchSplitDimExpansion(srcVecTy.getShape(), resShape,
+                                     splitDimGroups))
+    return consumerLayout;
+
+  SmallVector<int64_t> laneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
+  SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
+  if (laneLayout.empty())
----------------
Jianhui-Li wrote:

Good catch. Should be removed. 

https://github.com/llvm/llvm-project/pull/227560