[llvm-branch-commits] [llvm] [CodeGen] Add initial multi-def rematerialization support (PR #197580)

[Matt Arsenault via llvm-branch-commits]

================
@@ -158,27 +159,62 @@ void Rematerializer::transferUserImpl(RegisterIdx FromRegIdx,
   assert(FromRegIdx != ToRegIdx && "identical registers");
   assert(getOriginOrSelf(FromRegIdx) == getOriginOrSelf(ToRegIdx) &&
          "unrelated registers");
+  assert(any_of(getReg(FromRegIdx).Uses,
+                [&](const auto RegionUsers) {
+                  return RegionUsers.second.contains(&UserMI);
+                }) &&
+         "not a user");
 
   LLVM_DEBUG(dbgs() << "User transfer from " << printID(FromRegIdx) << " to "
                     << printID(ToRegIdx) << ": " << printUser(&UserMI) << '\n');
 
-  UserMI.substituteRegister(getReg(FromRegIdx).getDefReg(),
-                            getReg(ToRegIdx).getDefReg(), 0, TRI);
+  Register FromReg = getReg(FromRegIdx).getDefReg(),
+           ToReg = getReg(ToRegIdx).getDefReg();
+  UserMI.substituteRegister(FromReg, ToReg, 0, TRI);
   LISUpdates.insert(FromRegIdx);
   LISUpdates.insert(ToRegIdx);
 
-  // If the user is rematerializable, we must change its dependency to the
-  // new register.
-  if (RegisterIdx UserRegIdx = getDefRegIdx(UserMI); UserRegIdx != NoReg) {
-    // Look for the user's dependency that matches the register.
-    for (RegisterIdx &DepRegIdx : Regs[UserRegIdx].Dependencies) {
-      if (DepRegIdx == FromRegIdx) {
-        DepRegIdx = ToRegIdx;
+  RegisterIdx UserRegIdx = getDefRegIdx(UserMI);
+  if (UserRegIdx == NoReg)
+    return;
+
+  // When the user is rematerializable, we must reflect the change in its
+  // dependencies.
+  Reg &UserReg = Regs[UserRegIdx];
+  SmallVectorImpl<RegisterIdx> &UserDeps = Regs[UserRegIdx].Dependencies;
+  bool IsNewDep = true;
+  if (UserReg.Defs.size() > 1) {
+    // Other defining MIs might already be using the new register.
+    IsNewDep = find(UserDeps, ToRegIdx) == UserDeps.end();
----------------
arsenm wrote:

is_contained 

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