[Mlir-commits] [mlir] [MLIR][XeGPU] Add 2D `vector.multi_reduction` optimization (PR #171154)

[Artem Kroviakov]

================
@@ -472,6 +562,17 @@ struct XeGPUOptimizeBlockLoadsPass final
           auto laneData = layout.getEffectiveLaneDataAsInt();
           return !canBeOptimizedForTranspose(laneLayout, laneData);
         });
+
+    target.addDynamicallyLegalOp<vector::MultiDimReductionOp>(
+        [=](Operation *op) -> bool {
+          auto layout = xegpu::getDistributeLayoutAttr(op->getResult(0));
+          if (!layout || !layout.isForSubgroup())
----------------
akroviakov wrote:

Propagation can fail or some intermediate transform creates an op without setting a layout. We indicate a reduction op without sg layout as an already legal op, because no pattern in this pass applies to such a reduction.

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