[AArch64] Try to fold lsl + ldr/str

[jaykang10]

Summary

On ISelDAG level, if there are multiple uses of lsl, the ldr/str, which are uses of the lsl, are not folded.
Additionaly, if the lsl is not in same block with ldr/str, they are not folded too.
ldr/str can be folded with lsl 2 for 32 bit and lsl 3 for 64 bit as below.

%7:gpr64 = UBFMXri killed %6:gpr64, 62, 61
%8:gpr32 = LDRWroX %1:gpr64common, %7:gpr64, 0, 0
STRWroX killed %9:gpr32, %1:gpr64common, %7:gpr64, 0, 0
==>
%8:gpr32 = LDRWroX %1:gpr64common, %6:gpr64, 0, 1
STRWroX killed %9:gpr32, %1:gpr64common, %6:gpr64, 0, 1

[dmgreen]

Could this just update the code in https://github.com/llvm/llvm-project/blob/1f5510326375cb9a7c747aab26080699f7952d74/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp#L510?

It sounds like FeatureLSLFast could be more generally enabled, depending on what the conditions are for when this is fast.

[jaykang10]

In D115029#3169432, @dmgreen wrote:

Could this just update the code in https://github.com/llvm/llvm-project/blob/1f5510326375cb9a7c747aab26080699f7952d74/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp#L510?

Ah, you are right! It looks the isWorthFoldingSHL is doing same thing and it will work within a block.
However, if the lsl is in different block with the ldr/str, it would not work because there would be CopyFromReg for lsl in the block of ldr/str.

It sounds like FeatureLSLFast could be more generally enabled, depending on what the conditions are for when this is fast.

Yep, I agree with you. At least, we need to enable it for neoverse.

llvm/lib/Target/AArch64/AArch64MIPeepholeOpt.cpp

 //
 //    %3:gpr32 = ORRWrs $wzr, %2, 0
 //    %4:gpr64 = SUBREG_TO_REG 0, %3, %subreg.sub_32
 //
 //    If AArch64's 32-bit form of instruction defines the source operand of
 //    ORRWrs, we can remove the ORRWrs because the upper 32 bits of the source
 //    operand are set to zero.
 //
+// 3. Try to fold lsl + ldr/str
+//
+//    ldr/str can be folded with `lsl 2` for 32 bit and `lsl 3` for 64 bit as
+//    below.
+//
+//    %7:gpr64 = UBFMXri killed %6:gpr64, 62, 61
+//    %8:gpr32 = LDRWroX %1:gpr64common, %7:gpr64, 0, 0
+//    STRWroX killed %9:gpr32, %1:gpr64common, %7:gpr64, 0, 0
+//    ==>
+//    %8:gpr32 = LDRWroX %1:gpr64common, %6:gpr64, 0, 1
+//    STRWroX killed %9:gpr32, %1:gpr64common, %6:gpr64, 0, 1
+//
 //===----------------------------------------------------------------------===//
 
 #include "AArch64ExpandImm.h"
 #include "AArch64InstrInfo.h"
 #include "MCTargetDesc/AArch64AddressingModes.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
struct AArch64MIPeepholeOpt : public MachineFunctionPass {
   MachineLoopInfo *MLI;
   MachineRegisterInfo *MRI;
 
   template <typename T>
   bool visitAND(MachineInstr &MI,
                 SmallSetVector<MachineInstr *, 8> &ToBeRemoved);
   bool visitORR(MachineInstr &MI,
                 SmallSetVector<MachineInstr *, 8> &ToBeRemoved);
+  bool visitLDRSTR(MachineInstr &MI,
+                   SmallSetVector<MachineInstr *, 8> &ToBeRemoved,
+                   bool is32Bit);
   bool runOnMachineFunction(MachineFunction &MF) override;
 
   StringRef getPassName() const override {
     return "AArch64 MI Peephole Optimization pass";
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
bool AArch64MIPeepholeOpt::visitORR(
   MI.getOperand(0).setReg(DefReg);
   ToBeRemoved.insert(&MI);
 
   LLVM_DEBUG({ dbgs() << "Removed: " << MI << "\n"; });
 
   return true;
 }
 
+bool AArch64MIPeepholeOpt::visitLDRSTR(
+    MachineInstr &MI, SmallSetVector<MachineInstr *, 8> &ToBeRemoved,
+    bool is32Bit) {
+  Register SrcReg = MI.getOperand(2).getReg();
+  MachineInstr *SrcMI = MRI->getUniqueVRegDef(SrcReg);
+  if (!SrcMI)
+    return false;
+
+  if (SrcMI->getOpcode() != AArch64::UBFMXri)
+    return false;
+
+  unsigned Immr = SrcMI->getOperand(2).getImm();
+  unsigned Imms = SrcMI->getOperand(3).getImm();
+  // Check lsl 2 for 32bit ldr/str.
+  if (is32Bit && (Immr != 62 || Imms != 61))
+    return false;
+
+  // Check lsl 3 for 64bit ldr/str.
+  if (!is32Bit && (Immr != 61 || Imms != 60))
+    return false;
+
+  // Update ldr/str with lsl.
+  MI.getOperand(2).setReg(SrcMI->getOperand(1).getReg());
+  MI.getOperand(3).setImm(0);
+  MI.getOperand(4).setImm(1);
+
+  if (MRI->use_empty(SrcReg))
+    ToBeRemoved.insert(SrcMI);
+
+  LLVM_DEBUG({ dbgs() << "Removed: " << MI << "\n"; });
+
+  return true;
+}
+
 bool AArch64MIPeepholeOpt::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()))
     return false;
 
   TII = static_cast<const AArch64InstrInfo *>(MF.getSubtarget().getInstrInfo());
   MLI = &getAnalysis<MachineLoopInfo>();
   MRI = &MF.getRegInfo();
 
for (MachineInstr &MI : MBB) {
       case AArch64::ANDWrr:
         Changed = visitAND<uint32_t>(MI, ToBeRemoved);
         break;
       case AArch64::ANDXrr:
         Changed = visitAND<uint64_t>(MI, ToBeRemoved);
         break;
       case AArch64::ORRWrs:
         Changed = visitORR(MI, ToBeRemoved);
+        break;
+      case AArch64::STRWroX:
+      case AArch64::LDRWroX:
+        Changed = visitLDRSTR(MI, ToBeRemoved, true /* is32Bit */);
+        break;
+      case AArch64::STRXroX:
+      case AArch64::LDRXroX:
+        Changed = visitLDRSTR(MI, ToBeRemoved, false /* is32Bit */);
+        break;
       }
     }
   }
 
   for (MachineInstr *MI : ToBeRemoved)
     MI->eraseFromParent();
 
   return Changed;
 }
 
 FunctionPass *llvm::createAArch64MIPeepholeOptPass() {
   return new AArch64MIPeepholeOpt();
 }

llvm/test/CodeGen/AArch64/aarch64-split-and-bitmask-immediate.ll

 ; CHECK-NEXT:    cmp x8, x3
 ; CHECK-NEXT:    b.gt .LBB7_4
 ; CHECK-NEXT:  .LBB7_2: // %loop
 ; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
 ; CHECK-NEXT:    cmp x8, x9
 ; CHECK-NEXT:    b.hs .LBB7_1
 ; CHECK-NEXT:  // %bb.3: // %if.then
 ; CHECK-NEXT:    // in Loop: Header=BB7_2 Depth=1
-; CHECK-NEXT:    lsl x10, x8, #3
-; CHECK-NEXT:    ldr x11, [x1, x10]
-; CHECK-NEXT:    str x11, [x2, x10]
+; CHECK-NEXT:    ldr x10, [x1, x8, lsl #3]
+; CHECK-NEXT:    str x10, [x2, x8, lsl #3]
 ; CHECK-NEXT:    b .LBB7_1
 ; CHECK-NEXT:  .LBB7_4: // %exit
 ; CHECK-NEXT:    ret
 loop.ph:
   br label %loop
 
 loop:
   %iv = phi i64 [ %inc, %for.inc ], [ 0, %loop.ph ]

llvm/test/CodeGen/AArch64/extract-bits.ll

 ; ---------------------------------------------------------------------------- ;
 
 ; https://bugs.llvm.org/show_bug.cgi?id=38938
 define void @pr38938(i32* %a0, i64* %a1) nounwind {
 ; CHECK-LABEL: pr38938:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ldr x8, [x1]
 ; CHECK-NEXT:    ubfx x8, x8, #21, #10
-; CHECK-NEXT:    lsl x8, x8, #2
-; CHECK-NEXT:    ldr w9, [x0, x8]
+; CHECK-NEXT:    ldr w9, [x0, x8, lsl #2]
 ; CHECK-NEXT:    add w9, w9, #1
-; CHECK-NEXT:    str w9, [x0, x8]
+; CHECK-NEXT:    str w9, [x0, x8, lsl #2]
 ; CHECK-NEXT:    ret
   %tmp = load i64, i64* %a1, align 8
   %tmp1 = lshr i64 %tmp, 21
   %tmp2 = and i64 %tmp1, 1023
   %tmp3 = getelementptr inbounds i32, i32* %a0, i64 %tmp2
   %tmp4 = load i32, i32* %tmp3, align 4
   %tmp5 = add nsw i32 %tmp4, 1
   store i32 %tmp5, i32* %tmp3, align 4