[llvm] [AMDGPU] Rematerialize scalar loads (PR #68778)

[Jay Foad via llvm-commits]

================
@@ -2434,6 +2445,103 @@ bool SIInstrInfo::expandPostRAPseudo(MachineInstr &MI) const {
   return true;
 }
 
+void SIInstrInfo::reMaterialize(MachineBasicBlock &MBB,
+                                MachineBasicBlock::iterator I, Register DestReg,
+                                unsigned SubIdx, const MachineInstr &Orig,
+                                const TargetRegisterInfo &RI) const {
+
+  // Try shrinking the instruction to remat only the part needed for current
+  // context.
+  // TODO: Handle more cases.
+  unsigned Opcode = Orig.getOpcode();
+  switch (Opcode) {
+  case AMDGPU::S_LOAD_DWORDX16_IMM:
+  case AMDGPU::S_LOAD_DWORDX8_IMM: {
+    if (SubIdx != 0)
+      break;
+
+    if (I == MBB.end())
+      break;
+
+    if (I->isBundled())
+      break;
+
+    // Look for a single use of the register that is also a subreg.
+    Register RegToFind = Orig.getOperand(0).getReg();
+    int SingleUseIdx = -1;
+    for (unsigned i = 0, e = I->getNumOperands(); i != e; ++i) {
+      const MachineOperand &CandMO = I->getOperand(i);
+      if (!CandMO.isReg())
+        continue;
+      Register CandReg = CandMO.getReg();
+      if (!CandReg)
+        continue;
+
+      if (CandReg == RegToFind) {
+        if (SingleUseIdx == -1 && CandMO.isUse()) {
+          SingleUseIdx = i;
+        } else {
+          SingleUseIdx = -1;
+          break;
+        }
+      }
+    }
+    if (SingleUseIdx == -1)
+      break;
+    MachineOperand &UseMO = I->getOperand(SingleUseIdx);
+    if (UseMO.getSubReg() == AMDGPU::NoSubRegister)
+      break;
----------------
jayfoad wrote:

I think you can ignore def operands. And perhaps the whole loop can be a bit simpler, something like:
```suggestion
    MachineOperand *UseMO = nullptr;
    for (auto &CandMO : I->operands()) {
      if (!CandMO.isReg() || CandMO.getReg() != RegToFind || CandMO.isDef())
        continue;
      if (UseMO) {
        UseMO = nullptr;
        break;
      }
      UseMO = &CandMO;
    }
    if (!UseMO || UseMO.getSubReg() == AMDGPU::NoSubRegister)
      break;
```

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