[llvm] Re-apply "[AMDGPU][Scheduler] Use MIR-level rematerializer in rematerialization stage (#189491)" (PR #192443)

[Lucas Ramirez via llvm-commits]

================
@@ -1594,44 +1591,52 @@ bool PreRARematStage::initGCNSchedStage() {
       // in at least one target region.
       if (!Cand.maybeBeneficial(TargetRegions, RPTargets)) {
         REMAT_DEBUG(dbgs() << "Interrupt round on stale score for "
-                           << Cand.print() << " | " << *Cand.Remat->DefMI);
+                           << Cand.print() << " | "
+                           << Remater.printRematReg(Cand.RegIdx));
         break;
       }
       CandidateOrder.pop_back();
-      RematReg &Remat = *Cand.Remat;
-
-      // Remove the register from all regions where it is a live-in or live-out
-      // and rematerialize it.
-      REMAT_DEBUG(dbgs() << "** REMAT " << PrintCandidate(Cand) << '\n');
-      removeFromLiveMaps(Remat.getReg(), Cand.LiveIn, Cand.LiveOut);
-      MachineInstr *RematMI = Cand.rematerialize(DAG);
-
-      // Every rematerialization we do here is likely to move the instruction
-      // into a higher frequency region, increasing the total sum latency of the
-      // instruction itself. This is acceptable if we are eliminating a spill in
-      // the process, but when the goal is increasing occupancy we get nothing
-      // out of rematerialization if occupancy is not increased in the end; in
-      // such cases we want to roll back the rematerialization.
-      if (TargetOcc) {
-        RollbackInfo &Rollback =
-            Rollbacks.emplace_back(&Remat, Cand.LiveIn, Cand.LiveOut);
-        Rollback.RematMI = RematMI;
-        // Make the original MI a debug value so that it does not influence
-        // scheduling and replace all read registers with a sentinel register to
-        // prevent operands to appear in use-lists of other MIs during LIS
-        // updates. Store mappings between operand indices and original
-        // registers for potential rollback.
-        Remat.DefMI->setDesc(DAG.TII->get(TargetOpcode::DBG_VALUE));
-        for (auto [Idx, MO] : enumerate(Remat.DefMI->operands())) {
-          if (MO.isReg() && MO.readsReg()) {
-            Rollback.RegMap.insert({Idx, MO.getReg()});
-            MO.setReg(Register());
-          }
+
+#ifdef EXPENSIVE_CHECKS
+      // All uses are known to be available / live at the remat point. Thus,
+      // the uses should already be live in to the using region.
+      for (MachineOperand &MO : Reg.DefMI->operands()) {
+        if (!MO.isReg() || !MO.getReg() || !MO.readsReg())
+          continue;
+
+        Register UseReg = MO.getReg();
+        if (!UseReg.isVirtual())
+          continue;
+
+        LiveInterval &LI = DAG.LIS->getInterval(UseReg);
+        LaneBitmask LM = DAG.MRI.getMaxLaneMaskForVReg(MO.getReg());
+        if (LI.hasSubRanges() && MO.getSubReg())
+          LM = DAG.TRI->getSubRegIndexLaneMask(MO.getSubReg());
+
+        const unsigned UseRegion = Reg.Uses.begin()->first;
+        LaneBitmask LiveInMask = DAG.LiveIns[UseRegion].at(UseReg);
+        LaneBitmask UncoveredLanes = LM & ~(LiveInMask & LM);
+        // If this register has lanes not covered by the LiveIns, be sure they
+        // do not map to any subrange. ref:
+        // machine-scheduler-sink-trivial-remats.mir::omitted_subrange
+        if (UncoveredLanes.any()) {
+          assert(LI.hasSubRanges());
+          for (LiveInterval::SubRange &SR : LI.subranges())
+            assert((SR.LaneMask & UncoveredLanes).none());
         }
-      } else {
-        // Just delete the original instruction if it cannot be rolled back.
-        DAG.deleteMI(Remat.DefRegion, Remat.DefMI);
       }
+#endif
----------------
lucas-rami wrote:

This is the block that moved from `PreRARematStage::ScoredRemat::rematerialize`, where it could not read e.g., the `DAG`.

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