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

[Kazu Hirata]

================
@@ -8692,6 +8694,108 @@ DeinterleaveOp::getIterationShapeToUnroll() {
   return llvm::to_vector<4>(getResultVectorType().getShape());
 }
 
+//===----------------------------------------------------------------------===//
+// VectorDialect canonicalization
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// 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();
+        if (!replacements.contains(useRegion)) {
+          replacements[useRegion] = Value();
----------------
kazutakahirata wrote:

May I suggest using `try_emplace` here to avoid performing up to four separate lookups on `replacements` for `useRegion`?

```cpp
        auto [it, inserted] = replacements.try_emplace(useRegion);
        if (inserted) {
          for (Value operand : parent->getOperands()) {
            ...
            it->second = replacement;
            break;
          }
        }
        if (Value replacement = it->second) {
```

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