[Mlir-commits] [mlir] [MLIR][XeGPU] Wrap layout with a slice attr when propagating broadcast (PR #169054)

[Igor Zamyatin]

================
@@ -593,10 +597,19 @@ void LayoutInfoPropagation::visitShapeCastOp(
     return;
   }
   int64_t slicedDim = resultTy.getShape()[0] == 1 ? 0 : 1;
-  xegpu::SliceAttr sliceLayout = xegpu::SliceAttr::get(
-      shapeCast->getContext(), cast<xegpu::LayoutAttr>(resultLayout.get()),
-      DenseI64ArrayAttr::get(shapeCast->getContext(), {slicedDim}));
-  propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(sliceLayout)));
+  LayoutInfo operandLayout;
+  if (auto sliceResultAttr = dyn_cast<xegpu::SliceAttr>(resultLayout.get())) {
+    auto sliceDims = sliceResultAttr.getDims().asArrayRef();
+    if (sliceDims.size() == 1 && sliceDims[0] == slicedDim)
+      operandLayout = resultLayout;
----------------
Garra1980 wrote:

maybe assert on else condition?

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