[Mlir-commits] [mlir] [MLIR][XeGPU] Recover temporary layout from Anchor Layout (PR #191947)

[Charitha Saumya]

================
@@ -1187,33 +1360,41 @@ xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
   // For vector::InsertStridedSliceOp, infer source layout from result layout.
   // Dest vector must have the same layout as the result.
   if (auto insertSlice = dyn_cast<vector::InsertStridedSliceOp>(op)) {
-    if (!resLayout)
-      return xegpu::DistributeLayoutAttr();
-    if (idx == 0)
+    if (idx == 0) {
       return xegpu::inferInsertStridedSliceSourceLayout(
           resLayout, insertSlice.getDestVectorType().getShape(),
           insertSlice.getSourceVectorType().getShape());
+    }
     if (idx == 1)
       return resLayout;
   }
 
   // For vector::TransposeOp, infer source layout from result layout using
   // permutation.
   if (auto transpose = dyn_cast<vector::TransposeOp>(op)) {
-    if (!resLayout)
-      return xegpu::DistributeLayoutAttr();
     return xegpu::inferTransposeSourceLayout(resLayout,
                                              transpose.getPermutation());
   }
 
-  // For elementwise operations, all operands must have the same layout as the
-  // result.
-  if (OpTrait::hasElementwiseMappableTraits(op) && op->getNumResults() == 1) {
-    if (!resLayout)
-      return xegpu::DistributeLayoutAttr();
+  if (isa<VectorType>(operand.get().getType()) &&
+      !dyn_cast<xegpu::AnchorLayoutInterface>(op)) {
+    // For elementwise operations, all operands must have the same layout as the
+    // result.
+    // if (OpTrait::hasElementwiseMappableTraits(op) && op->getNumResults() ==
+    // 1) {
----------------
charithaintc wrote:

cleanup the comment. I think this is missing some sanity checks. Are we handling all possible cases? worth adding checks to report unhandled cases. 

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