[llvm] [AArch64] Fold UMOV(lane 0) + GPR store in FPR store (PR #199139)

[David Green via llvm-commits]

================
@@ -3012,6 +3022,140 @@ bool AArch64LoadStoreOpt::tryToMergeIndexLdSt(MachineBasicBlock::iterator &MBBI,
   return false;
 }
 
+// Map a GPR store opcode to its FPR equivalent at the same data width.
+// Returns 0 if no mapping exists.
+static unsigned getGPRToFPRStoreOpcode(unsigned GPRStoreOpc) {
+  switch (GPRStoreOpc) {
+  // Unsigned immediate.
+  case AArch64::STRBBui:
+    return AArch64::STRBui;
+  case AArch64::STRHHui:
+    return AArch64::STRHui;
+  case AArch64::STRWui:
+    return AArch64::STRSui;
+  case AArch64::STRXui:
+    return AArch64::STRDui;
+  // Unscaled immediate.
+  case AArch64::STURBBi:
+    return AArch64::STURBi;
+  case AArch64::STURHHi:
+    return AArch64::STURHi;
+  case AArch64::STURWi:
+    return AArch64::STURSi;
+  case AArch64::STURXi:
+    return AArch64::STURDi;
+  // Register offset.
+  case AArch64::STRBBroW:
+    return AArch64::STRBroW;
+  case AArch64::STRBBroX:
+    return AArch64::STRBroX;
+  case AArch64::STRHHroW:
+    return AArch64::STRHroW;
+  case AArch64::STRHHroX:
+    return AArch64::STRHroX;
+  case AArch64::STRWroW:
+    return AArch64::STRSroW;
+  case AArch64::STRWroX:
+    return AArch64::STRSroX;
+  case AArch64::STRXroW:
+    return AArch64::STRDroW;
+  case AArch64::STRXroX:
+    return AArch64::STRDroX;
+  default:
+    return 0;
+  }
+}
+
+// Given a UMOV-lane-0 opcode, return the sub-register index to extract from
+// the vector register, or 0 if the opcode is not a supported UMOV.
+static unsigned getUMOVSubRegIdx(unsigned UMOVOpc) {
+  switch (UMOVOpc) {
+  case AArch64::UMOVvi8_idx0:
+    return AArch64::bsub;
+  case AArch64::UMOVvi16_idx0:
+    return AArch64::hsub;
+  case AArch64::UMOVvi32_idx0:
+    return AArch64::ssub;
+  case AArch64::UMOVvi64_idx0:
+    return AArch64::dsub;
+  default:
+    return 0;
+  }
+}
+
+bool AArch64LoadStoreOpt::tryToReplaceUMOVStore(
+    MachineBasicBlock::iterator &MBBI) {
+  MachineInstr &StoreMI = *MBBI;
+
+  unsigned FPRStoreOpc = getGPRToFPRStoreOpcode(StoreMI.getOpcode());
+  if (!FPRStoreOpc)
+    return false;
+
+  MachineBasicBlock *MBB = StoreMI.getParent();
+  MCPhysReg StoreValReg = StoreMI.getOperand(0).getReg();
+
+  if (!StoreMI.getOperand(0).isKill())
+    return false;
+
+  // Scan backward to find the UMOV that defines the store's value register.
+  MachineInstr *UMOVMI = nullptr;
+  MachineBasicBlock::iterator B = MBB->begin();
+  unsigned SubRegIdx = 0;
+  unsigned Count = 0;
+  for (auto It = MBBI; It != B && !UMOVMI;) {
+    MachineInstr &MI = *--It;
+    if (MI.isDebugInstr())
+      continue;
+    if (++Count > UMOVFoldLimit)
+      return false;
+    if (MI.readsRegister(StoreValReg, TRI))
+      return false;
+    if (MI.definesRegister(StoreValReg, TRI)) {
+      SubRegIdx = getUMOVSubRegIdx(MI.getOpcode());
+      if (!SubRegIdx)
+        return false;
+      UMOVMI = &MI;
----------------
davemgreen wrote:

break here?

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