[Mlir-commits] [mlir] [mlir][vector] Refine select uses on uniform reduction paths (PR #235640)

[Federico Bruzzone]

================
@@ -8713,11 +8715,102 @@ struct SelectBroadcastCondition : OpRewritePattern<arith::SelectOp> {
     return success();
   }
 };
+
+/// Conservatively compute the regions reachable from an operation entry.
+/// Region-to-region edges are included: an entry-only check is insufficient for
+/// loops and operations that enter one region through another.
+static bool isRegionReachable(RegionBranchOpInterface branch, Region *target,
+                              ArrayRef<Attribute> operands) {
+  SmallVector<RegionSuccessor> successors;
+  branch.getEntrySuccessorRegions(operands, successors);
+  SmallVector<Region *> worklist;
+  llvm::SmallPtrSet<Region *, 4> visited;
+  auto append = [&] {
+    for (RegionSuccessor successor : successors)
+      if (Region *region = successor.getSuccessor())
+        if (region->getParentOp() == branch.getOperation() &&
+            visited.insert(region).second)
+          worklist.push_back(region);
+  };
+  append();
+  while (!worklist.empty()) {
+    Region *region = worklist.pop_back_val();
+    if (region == target)
+      return true;
+    successors.clear();
+    branch.getSuccessorRegions(*region, successors);
+    append();
+  }
+  return false;
+}
+
+/// Infer a reduction's value only when the opposite value cannot reach the use.
+struct SelectReductionBranch : OpRewritePattern<arith::SelectOp> {
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(arith::SelectOp select,
+                                PatternRewriter &rewriter) const override {
+    if (!isa<VectorType>(select.getCondition().getType()))
+      return failure();
+    DenseMap<Region *, Value> replacements;
+    bool changed = false;
+    for (OpOperand &use :
+         llvm::make_early_inc_range(select.getResult().getUses())) {
+      for (Operation *child = use.getOwner();
+           auto *parent = child->getParentOp(); child = parent) {
+        auto branch = dyn_cast<RegionBranchOpInterface>(parent);
+        if (!branch)
+          continue;
+        Region *useRegion = child->getParentRegion();
+        auto [it, inserted] = replacements.try_emplace(useRegion);
+        if (inserted) {
+          for (Value operand : parent->getOperands()) {
+            auto reduction = operand.getDefiningOp<ReductionOp>();
+            if (!reduction ||
+                !reduction.getDest().getType().isSignlessInteger(1) ||
----------------
FedericoBruzzone wrote:

This looks redundant...

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