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

[Matt Arsenault via llvm-branch-commits]

================
@@ -88,47 +91,59 @@ namespace llvm {
 /// In its nomenclature, the rematerializer differentiates between "original
 /// registers" (registers that were present when it analyzed the function) and
 /// rematerializations of these original registers. Rematerializations have an
-/// "origin" which is the index of the original regiser they were rematerialized
-/// from (transitivity applies; a rematerialization and all of its own
-/// rematerializations have the same origin). Semantically, only original
-/// registers have rematerializations.
+/// "origin" which is the index of the original register they were
+/// rematerialized from (transitivity applies; a rematerialization and all of
+/// its own rematerializations have the same origin). Semantically, only
+/// original registers have rematerializations.
+///
+/// Dealing with sub-registers is complicated, we have to handle dead-defs,
+/// undef flags, and connected components
 class Rematerializer {
 public:
   /// Index type for rematerializable registers.
   using RegisterIdx = unsigned;
 
-  /// A rematerializable register defined by a single machine instruction.
+  /// A rematerializable register, potentially defined by multiple instructions.
   ///
   /// A rematerializable register has a set of dependencies, which correspond
-  /// to the unique read register operands of its defining instruction and which
-  /// can themselves be rematerializable. Operand indices corresponding to
-  /// unrematerializable dependencies are managed by and queried from the
-  /// rematerializer, whereas rematerializable ones are part of this struct and
-  /// identified through their register index.
+  /// to the unique read register operands of its defining instruction(s) and
+  /// which can themselves be rematerializable. Operands of defining
+  /// instructions corresponding to unrematerializable dependencies are managed
+  /// by and queried from the rematerializer, whereas rematerializable ones are
+  /// part of this struct and identified through their register index.
   ///
   /// A rematerializable register also has an arbitrary number of users in an
   /// arbitrary number of regions, potentially including its own defining
   /// region. When rematerializations lead to operand changes in users, a
   /// register may find itself without any user left, at which point the
-  /// rematerializer deletes it (setting its defining MI to nullptr).
+  /// rematerializer deletes it (emptying \ref Reg::Defs).
   struct Reg {
-    /// Single MI defining the rematerializable register.
-    MachineInstr *DefMI;
-    /// Defining region of \p DefMI.
+    /// All instructions that define the register, in program order.
+    SmallVector<MachineInstr *, 1> Defs;
+    /// Defining region of the register.
     unsigned DefRegion;
     /// The rematerializable register's lane bitmask.
     LaneBitmask Mask;
 
     using RegionUsers = SmallDenseSet<MachineInstr *, 4>;
-    /// Uses of the register, mapped by region.
+    /// Uses of the register, mapped by region. Users that also define a part of
+    /// the register are considered defs and not accounted for here.
     SmallDenseMap<unsigned, RegionUsers, 2> Uses;
+
     /// This register's rematerializable dependencies, one per unique
-    /// rematerializable register operand.
+    /// rematerializable register operand over all definitions.
     SmallVector<RegisterIdx, 2> Dependencies;
 
-    /// Returns the rematerializable register from its defining instruction.
+    MachineInstr *getFirstDef() const { return Defs.front(); }
+    MachineInstr *getLastDef() const { return Defs.back(); }
+
+    /// Returns the rematerializable register from one of its defining
+    /// instructions.
     Register getDefReg() const {
-      assert(DefMI && "defining instruction was deleted");
+      const MachineInstr *DefMI = getFirstDef();
+      assert(DefMI && "defining instruction(s) were deleted");
+      if (!DefMI->getOperand(0).isDef())
+        dbgs() << *DefMI;
       assert(DefMI->getOperand(0).isDef() && "not a register def");
----------------
arsenm wrote:

```suggestion
      assert(DefMI && DefMI->getOperand(0).isDef() && "not a register def");
```

Stray debug printing 

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