[llvm] [InstCombine] Extend folding of aggregate construction to cases when source aggregates are partially available (PR #100828)

[Nikita Popov via llvm-commits]

================
@@ -1117,9 +1124,68 @@ Instruction *InstCombinerImpl::foldAggregateConstructionIntoAggregateReuse(
     // aggregate produced by OrigIVI must have been originally extracted from
     // the same aggregate. Is that so? Can we find said original aggregate?
     SourceAggregate = FindCommonSourceAggregate(UseBB, Pred);
-    if (Describe(SourceAggregate) != AggregateDescription::Found)
-      return nullptr; // Give up.
-    IV.first->second = *SourceAggregate;
+    if (Describe(SourceAggregate) == AggregateDescription::Found) {
+      FoundSrcAgg = true;
+      IV.first->second = *SourceAggregate;
+    } else {
+      // If UseBB is the single successor of Pred, we can add InsertValue to
+      // Pred.
+      auto *BI = dyn_cast<BranchInst>(Pred->getTerminator());
+      if (!BI || !BI->isUnconditional())
+        return nullptr;
+    }
+  }
+
+  if (!FoundSrcAgg)
+    return nullptr;
----------------
nikic wrote:

For example, here is a slight variant of test2 with the current patch:
```
define {ptr, i64} @test2(i1 %cond1, i1 %cond2, ptr %p1, ptr %p2) {
; CHECK-LABEL: define { ptr, i64 } @test2(
; CHECK-SAME: i1 [[COND1:%.*]], i1 [[COND2:%.*]], ptr [[P1:%.*]], ptr [[P2:%.*]]) {
; CHECK-NEXT:    br i1 [[COND1]], label %[[BBB1:.*]], label %[[BBB4:.*]]
; CHECK:       [[BBB1]]:
; CHECK-NEXT:    br i1 [[COND2]], label %[[BBB2:.*]], label %[[BBB3:.*]]
; CHECK:       [[BBB2]]:
; CHECK-NEXT:    [[CALL1:%.*]] = call { ptr, i64 } @bar(i64 0)
; CHECK-NEXT:    br label %[[EXIT:.*]]
; CHECK:       [[BBB3]]:
; CHECK-NEXT:    [[VAL31:%.*]] = load ptr, ptr [[P1]], align 8
; CHECK-NEXT:    [[VAL32:%.*]] = load i64, ptr [[P2]], align 4
; CHECK-NEXT:    [[TMP1:%.*]] = insertvalue { ptr, i64 } poison, ptr [[VAL31]], 0
; CHECK-NEXT:    [[TMP2:%.*]] = insertvalue { ptr, i64 } [[TMP1]], i64 [[VAL32]], 1
; CHECK-NEXT:    br label %[[EXIT]]
; CHECK:       [[BBB4]]:
; CHECK-NEXT:    [[VAL33:%.*]] = load ptr, ptr [[P1]], align 8
; CHECK-NEXT:    [[VAL34:%.*]] = load i64, ptr [[P2]], align 4
; CHECK-NEXT:    [[TMP3:%.*]] = insertvalue { ptr, i64 } poison, ptr [[VAL33]], 0
; CHECK-NEXT:    [[TMP4:%.*]] = insertvalue { ptr, i64 } [[TMP3]], i64 [[VAL34]], 1
; CHECK-NEXT:    br label %[[EXIT]]
; CHECK:       [[EXIT]]:
; CHECK-NEXT:    [[RES_MERGED:%.*]] = phi { ptr, i64 } [ [[CALL1]], %[[BBB2]] ], [ [[TMP2]], %[[BBB3]] ], [ [[TMP4]], %[[BBB4]] ]
; CHECK-NEXT:    ret { ptr, i64 } [[RES_MERGED]]
;
  br i1 %cond1, label %bbb1, label %bbb4

bbb1:
  br i1 %cond2, label %bbb2, label %bbb3

bbb2:
  %call1 = call {ptr, i64} @bar(i64 0)
  %val11 = extractvalue { ptr, i64 } %call1, 0
  %val12 = extractvalue { ptr, i64 } %call1, 1
  br label %exit

bbb3:
  %val21 = load ptr, ptr %p1
  %val22 = load i64, ptr %p2
  br label %exit

bbb4:
  %val31 = load ptr, ptr %p1
  %val32 = load i64, ptr %p2
  br label %exit

exit:
  %val1 = phi ptr [%val11, %bbb2], [%val21, %bbb3], [%val31, %bbb4]
  %val2 = phi i64 [%val12, %bbb2], [%val22, %bbb3], [%val32, %bbb4]
  %tmp = insertvalue { ptr, i64 } poison, ptr %val1, 0
  %res = insertvalue { ptr, i64 } %tmp, i64 %val2, 1
  ret {ptr, i64} %res
}
```

We remove extractvalues in one branch, but trade this with insertvalues in two branches. This doesn't look profitable, especially once we extend this to even more branches that don't fold.

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