[Mlir-commits] [mlir] [mlir][xegpu] Add support for `vector.reduction` and `vector.multi_reduction` subgroup to work-item distribution. (PR #180308)

[Artem Kroviakov]

================
@@ -551,8 +715,44 @@ void xegpu::populateXeGPUSgToWiDistributeTypeConversionAndLegality(
         }
         return !xegpu::getTemporaryLayout(dyn_cast<OpResult>(op->getResult(0)));
       });
+  // vector::ReductionOp is legal only if its source has no distribute layout
+  // attribute.
+  target.addDynamicallyLegalOp<vector::ReductionOp>(
+      [=](vector::ReductionOp op) -> bool {
+        auto layout = xegpu::getDistributeLayoutAttr(op.getVector());
+        return !layout;
+      });
+  // vector::MultiDimReductionOp op legality.
+  target.addDynamicallyLegalOp<vector::MultiDimReductionOp>(
+      [=](vector::MultiDimReductionOp op) -> bool {
+        // Check common conditions for subgroup multi reduction op.
+        if (!isValidSubgroupMultiReductionOp(op))
+          return true;
+        // Lane local reductions are illegal at this point and must be lowered.
+        return !isReductionLaneLocal(op);
----------------
akroviakov wrote:

Would be useful to comment in code or PR about the differences between the legality check here and `populateXeGPUSgToWiLowerVectorMultiReductionAndLegality` and why they need to be separated.

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