[llvm] [VPlan] Add vputils::reconstructSSA (PR #212209)

[Florian Hahn via llvm-commits]

================
@@ -1366,3 +1366,40 @@ void vputils::detail::pullOutPermutationsImpl(
     }
   }
 }
+
+// Implements the algorithm described in "Simple and Efficient Construction of
+// Static Single Assignment Form" by Braun et al.
+VPValue *vputils::reconstructSSA(VPBasicBlock *VPBB,
+                                 DenseMap<VPBasicBlock *, VPValue *> &Defs) {
+  assert(!Defs.empty() && "Defs shouldn't be empty");
+  assert(VPBB->getPlan() && "VPBB isn't reachable from entry");
+  if (VPValue *Def = Defs.lookup(VPBB))
+    return Def;
+  // If the entry block is reached and there's still no def, then Defs is
+  // missing a definition that covers this path.
+  assert(VPBB->getNumPredecessors() && "Not all paths have def");
+
+  if (VPBlockBase *Pred = VPBB->getSinglePredecessor())
+    return reconstructSSA(cast<VPBasicBlock>(Pred), Defs);
+
+  // Multiple predecessors, create a join.
+  Type *Ty = Defs.begin()->second->getScalarType();
+  auto *Phi = new VPPhi({}, VPIRFlags::getDefaultFlags(Instruction::PHI, Ty),
+                        DebugLoc::getUnknown(), "", Ty);
+  VPBB->insert(Phi, VPBB->getFirstNonPhi());
+  Defs[VPBB] = Phi;
+  for (auto *Pred : VPBB->predecessors())
+    Phi->addIncoming(reconstructSSA(cast<VPBasicBlock>(Pred), Defs));
+
+  // Fold away trivial phis.
+  // TODO: Remove phi users which have become trivial too.
+  if (all_equal(Phi->incoming_values())) {
+    VPValue *Common = Phi->getIncomingValue(0);
+    Phi->replaceAllUsesWith(Common);
----------------
fhahn wrote:

I think there may be cases where we add `Phi` to `Defs` for a different block in one of the recursions above, if in the recursion we fold another phi to `Phi`.

Then we may end up with a pointer to the deleted `Phi` in `Defs`, which may trigger use-after-free. Could we just rely on using regular VPlan DCE? Or remove all dead phis at the end?

Here's a test case that should show use-after-free with UBSan (created with assistance of AI)

```
TEST_F(VPUtilsTest, ReconstructSSAStaleDefsCrash) {
  VPlan &Plan = getPlan();
  VPBasicBlock *Y = Plan.getEntry();
  VPBasicBlock *Pj = Plan.createVPBasicBlock("");
  VPBasicBlock *Pk = Plan.createVPBasicBlock("");
  VPBasicBlock *B = Plan.createVPBasicBlock("");
  VPBasicBlock *X = Plan.createVPBasicBlock("");
  VPBasicBlock *Xb = Plan.createVPBasicBlock("");
  VPBasicBlock *T = Plan.createVPBasicBlock("");

  VPBlockUtils::connectBlocks(Y, Pj);
  VPBlockUtils::connectBlocks(X, Pj);
  VPBlockUtils::connectBlocks(Pj, B);
  VPBlockUtils::connectBlocks(Pj, Pk);
  VPBlockUtils::connectBlocks(Pk, B);
  VPBlockUtils::connectBlocks(B, X);
  VPBlockUtils::connectBlocks(B, Xb);
  VPBlockUtils::connectBlocks(Xb, X);
  VPBlockUtils::connectBlocks(B, T);
  VPBlockUtils::connectBlocks(X, T);

  VPValue *C = Plan.getConstantInt(32, 1);
  VPIRFlags AddFlags = VPIRFlags::getDefaultFlags(Instruction::Add);
  auto *Def = new VPInstruction(Instruction::Add, {C, C}, AddFlags);
  Y->appendRecipe(Def);

  DenseMap<VPBasicBlock *, VPValue *> Defs;
  Defs[Y] = Def;
  EXPECT_NE(vputils::reconstructSSA(T, Defs), nullptr);
}
```



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