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

[Jianhui Li]

================
@@ -473,6 +556,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())
+            return true;
+          if (auto reductionOp = dyn_cast<vector::MultiDimReductionOp>(op))
----------------
Jianhui-Li wrote:

nit: this check can be removed since it is always TRUE? 

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