[CodeGen] Fix incorrect rematerialization order in rematerializer (llvm#189485)

When rematerializing DAGs of registers wherein multiple paths exist
between some regsters of the DAG, it is possible that the
rematerialization determines an incorrect rematerialization order that
does not ensure that a register's dependencies are rematerialized before
itself; an invariant that is otherwise required.

This fixes that using a simpler recursive logic to determine a correct
rematerialization order that honors this invariant. A minimal unit test
is added that fails on the current implementation.

Author: lucas-rami

llvm/lib/CodeGen/Rematerializer.cpp

@@ -94,40 +94,26 @@ Rematerializer::rematerializeToPos(RegisterIdx RootIdx, unsigned UseRegion,
   assert(!DRI.DependencyMap.contains(RootIdx));
   LLVM_DEBUG(dbgs() << "Rematerializing " << printID(RootIdx) << '\n');
 
-  // Traverse the root's dependency DAG depth-first to find the set of
-  // registers we must rematerialize along with it and a legal order to
-  // rematerialize them in.
-  SmallVector<RegisterIdx, 4> DepDAG{RootIdx};
-  SmallSetVector<RegisterIdx, 8> RematOrder;
-  RematOrder.insert(RootIdx);
-  do {
-    RegisterIdx RegIdx = DepDAG.pop_back_val();
-    for (const Reg::Dependency &Dep : getReg(RegIdx).Dependencies) {
-      // The dependency may already have a rematerialization ready to use.
-      if (DRI.DependencyMap.contains(Dep.RegIdx))
-        continue;
-      // We may have already seen the dependency in the dependency DAG.
-      if (RematOrder.contains(Dep.RegIdx))
-        continue;
-      DepDAG.push_back(Dep.RegIdx);
-      RematOrder.insert(Dep.RegIdx);
-    }
-  } while (!DepDAG.empty());
-
-  // Rematerialize all necessary registers in the root's dependency DAG. At each
-  // rematerialization, dependencies should already be available.
-  RegisterIdx LastNewIdx;
-  for (RegisterIdx RegIdx : reverse(RematOrder)) {
-    assert(!DRI.DependencyMap.contains(RegIdx) && "useless remat");
-    SmallVector<Reg::Dependency, 2> Dependencies;
-    for (const Reg::Dependency &Dep : getReg(RegIdx).Dependencies)
-      Dependencies.emplace_back(Dep.MOIdx, DRI.DependencyMap.at(Dep.RegIdx));
-    LastNewIdx =
-        rematerializeReg(RegIdx, UseRegion, InsertPos, std::move(Dependencies));
-    DRI.DependencyMap.insert({RegIdx, LastNewIdx});
+  SmallVector<Reg::Dependency, 2> NewDeps;
+  // Copy all dependencies because recursive rematerialization of dependencies
+  // may invalidate references to the backing vector of registers.
+  SmallVector<Reg::Dependency, 2> OldDeps(getReg(RootIdx).Dependencies);
+  for (const Reg::Dependency &Dep : OldDeps) {
+    // Recursively rematerialize required dependencies at the same position as
+    // the root. Registers form a DAG so the recursion is guaranteed to
+    // terminate.
+    auto RematIdx = DRI.DependencyMap.find(Dep.RegIdx);
+    RegisterIdx NewDepRegIdx;
+    if (RematIdx == DRI.DependencyMap.end())
+      NewDepRegIdx = rematerializeToPos(Dep.RegIdx, UseRegion, InsertPos, DRI);
+    else
+      NewDepRegIdx = RematIdx->second;
+    NewDeps.emplace_back(Dep.MOIdx, NewDepRegIdx);
   }
-
-  return LastNewIdx;
+  RegisterIdx NewIdx =
+      rematerializeReg(RootIdx, UseRegion, InsertPos, std::move(NewDeps));
+  DRI.DependencyMap.insert({RootIdx, NewIdx});
+  return NewIdx;
 }
 
 void Rematerializer::transferUser(RegisterIdx FromRegIdx, RegisterIdx ToRegIdx,

llvm/unittests/CodeGen/RematerializerTest.cpp

@@ -318,6 +318,44 @@ body:             |
   EXPECT_TRUE(getMF().verify());
 }
 
+/// To rematerialize %3 along with all its dependencies before its only use in
+/// bb.1, we must first rematerialize %0 and %1 (in any order), then %2, and
+/// finally %3. The rematerializer had a rematerialization order bug wherein,
+/// because %0 is also used directly in the MI defining %3, it was
+/// rematerialized after %2, breaking the invariant that dependencies of a
+/// register must always be rematerialized before the register itself.
+TEST_F(RematerializerTest, MultiplePathsRematOrder) {
+  StringRef MIR = R"(
+name:            MultiplePathsRematOrder
+tracksRegLiveness: true
+machineFunctionInfo:
+  isEntryFunction: true
+body:             |
+  bb.0:
+    %0:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 0, implicit $exec, implicit $mode
+    %1:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 1, implicit $exec, implicit $mode
+    %2:vgpr_32 = V_ADD_U32_e32 %0, %1, implicit $exec
+    %3:vgpr_32 = V_ADD_U32_e32 %0, %2, implicit $exec
+  
+  bb.1:
+    S_NOP 0, implicit %3
+    S_ENDPGM 0
+...
+)";
+  ASSERT_TRUE(parseMIRAndInit(MIR, "MultiplePathsRematOrder"));
+  Rematerializer &Remater = getRematerializer();
+  Rematerializer::DependencyReuseInfo DRI;
+
+  const unsigned MBB1 = 1;
+  const RegisterIdx Add02 = 3;
+
+  // This call would previously fail.
+  Remater.rematerializeToRegion(Add02, MBB1, DRI);
+
+  Remater.updateLiveIntervals();
+  EXPECT_TRUE(getMF().verify());
+}
+
 /// Rematerializes a single register to multiple regions, tracking that
 /// rematerializations are linked correctly and making sure that the original
 /// register is deleted automatically when it no longer has any uses.