[RISCV] Teach RISCVMergeBaseOffset about cases where we use SHXADD to add some immediates.

[craig.topper]

Summary

For an addition with simm14 and simm15 immediates with 2 or 3 trailing bits,
we can use a shXadd instruction and an addi to do the addition.

This patch teaches RISCVMergeBaseOffset to see through this pattern.
I don't think the sh1add case occurs because we use two addis for that,
but I implemented it for completeness.

[sunshaoce]

llvm/test/CodeGen/RISCV/hoist-global-addr-base.ll

245

12848 -> 6424

[craig.topper]

Fix comment in test.

[craig.topper]

Add assertion that the ADDI immediate is a simm12.

[reames]

LGTM

[hctim]

Hi, looks like this patch broke the sanitizer buildbots (the fast builder, testing under ubsan). I can see it easily being missed as the bot was already red:

https://lab.llvm.org/buildbot/#/builders/5/builds/24595

Some instructions about helping reproduce sanitizer bots are here: https://github.com/google/sanitizers/wiki/SanitizerBotReproduceBuild

FAIL: LLVM :: CodeGen/RISCV/hoist-global-addr-base.ll (32134 of 66378)
******************** TEST 'LLVM :: CodeGen/RISCV/hoist-global-addr-base.ll' FAILED ********************
Script:
--
: 'RUN: at line 2';   /b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/llc -mtriple=riscv32  < /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/CodeGen/RISCV/hoist-global-addr-base.ll | /b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/FileCheck  /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/CodeGen/RISCV/hoist-global-addr-base.ll --check-prefixes=CHECK,RV32,RV32I
: 'RUN: at line 3';   /b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/llc -mtriple=riscv64  < /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/CodeGen/RISCV/hoist-global-addr-base.ll | /b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/FileCheck  /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/CodeGen/RISCV/hoist-global-addr-base.ll --check-prefixes=CHECK,RV64,RV64I
: 'RUN: at line 4';   /b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/llc -mtriple=riscv32 -mattr=+zba < /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/CodeGen/RISCV/hoist-global-addr-base.ll | /b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/FileCheck  /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/CodeGen/RISCV/hoist-global-addr-base.ll --check-prefixes=CHECK,RV32,RV32ZBA
: 'RUN: at line 5';   /b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/llc -mtriple=riscv64 -mattr=+zba < /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/CodeGen/RISCV/hoist-global-addr-base.ll | /b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/FileCheck  /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/CodeGen/RISCV/hoist-global-addr-base.ll --check-prefixes=CHECK,RV64,RV64ZBA
--
Exit Code: 2
Command Output (stderr):
--
/b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/Target/RISCV/RISCVMergeBaseOffset.cpp:239:30: runtime error: left shift of negative value -2048
    #0 0x5590c6f830fa in matchShiftedOffset /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/Target/RISCV/RISCVMergeBaseOffset.cpp:239:30
    #1 0x5590c6f830fa in detectAndFoldOffset /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/Target/RISCV/RISCVMergeBaseOffset.cpp:302:10
    #2 0x5590c6f830fa in (anonymous namespace)::RISCVMergeBaseOffsetOpt::runOnMachineFunction(llvm::MachineFunction&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/Target/RISCV/RISCVMergeBaseOffset.cpp:373:21
    #3 0x5590c7c61f9a in llvm::MachineFunctionPass::runOnFunction(llvm::Function&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/CodeGen/MachineFunctionPass.cpp:73:13
    #4 0x5590c8210919 in llvm::FPPassManager::runOnFunction(llvm::Function&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IR/LegacyPassManager.cpp:1430:27
    #5 0x5590c821940d in llvm::FPPassManager::runOnModule(llvm::Module&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IR/LegacyPassManager.cpp:1476:16
    #6 0x5590c821130c in runOnModule /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IR/LegacyPassManager.cpp:1545:27
    #7 0x5590c821130c in llvm::legacy::PassManagerImpl::run(llvm::Module&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IR/LegacyPassManager.cpp:535:44
    #8 0x5590c6055677 in compileModule(char**, llvm::LLVMContext&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/tools/llc/llc.cpp:730:8
    #9 0x5590c60530ab in main /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/tools/llc/llc.cpp:417:22
    #10 0x7f50f77ae09a in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x2409a) (BuildId: eb6a5dd378d22b1e695984462a799cd4c81cdc22)
    #11 0x5590c6027a59 in _start (/b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/llc+0x600ea59)
SUMMARY: UndefinedBehaviorSanitizer: undefined-behavior /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/Target/RISCV/RISCVMergeBaseOffset.cpp:239:30 in 
FileCheck error: '<stdin>' is empty.
FileCheck command line:  /b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/FileCheck /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/CodeGen/RISCV/hoist-global-addr-base.ll --check-prefixes=CHECK,RV32,RV32ZBA

llvm/lib/Target/RISCV/RISCVMergeBaseOffset.cpp

 //===----- RISCVMergeBaseOffset.cpp - Optimise address calculations  ------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
public:
   static char ID;
   bool runOnMachineFunction(MachineFunction &Fn) override;
   bool detectLuiAddiGlobal(MachineInstr &LUI, MachineInstr *&ADDI);
 
   bool detectAndFoldOffset(MachineInstr &HiLUI, MachineInstr &LoADDI);
   void foldOffset(MachineInstr &HiLUI, MachineInstr &LoADDI, MachineInstr &Tail,
                   int64_t Offset);
   bool matchLargeOffset(MachineInstr &TailAdd, Register GSReg, int64_t &Offset);
+  bool matchShiftedOffset(MachineInstr &TailShXAdd, Register GSReg,
+                          int64_t &Offset);
+
   RISCVMergeBaseOffsetOpt() : MachineFunctionPass(ID) {}
 
   MachineFunctionProperties getRequiredProperties() const override {
     return MachineFunctionProperties().set(
         MachineFunctionProperties::Property::IsSSA);
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
if (OffsetTail.getOpcode() == RISCV::ADDI ||
     LLVM_DEBUG(dbgs() << "  Offset Instr: " << OffsetTail);
     Offset = SignExtend64<32>(OffsetTail.getOperand(1).getImm() << 12);
     DeadInstrs.insert(&OffsetTail);
     return true;
   }
   return false;
 }
 
+// Detect patterns for offsets that are passed into a SHXADD instruction.
+// The offset has 1,2, or 3 trailing zeros and fits in simm13, simm14, simm15.
+// The constant is created with addi    voff, x0, C, and shXadd is used to
+// fill insert the trailing zeros and do the addition.
+//
+// HiLUI:      lui     vreg1, %hi(s)
+// LoADDI:     addi    vreg2, vreg1, %lo(s)
+// OffsetTail: addi    voff, x0, C
+// TailAdd:    shXadd  vreg4, voff, vreg2
+bool RISCVMergeBaseOffsetOpt::matchShiftedOffset(MachineInstr &TailShXAdd,
+                                                 Register GAReg,
+                                                 int64_t &Offset) {
+  assert((TailShXAdd.getOpcode() == RISCV::SH1ADD ||
+          TailShXAdd.getOpcode() == RISCV::SH2ADD ||
+          TailShXAdd.getOpcode() == RISCV::SH3ADD) &&
+         "Expected SHXADD instruction!");
+  // The first source is the shifted operand.
+  Register Rs1 = TailShXAdd.getOperand(1).getReg();
+
+  if (GAReg != TailShXAdd.getOperand(2).getReg())
+    return false;
+
+  // Can't fold if the register has more than one use.
+  if (!MRI->hasOneUse(Rs1))
+    return false;
+  // This can point to an ADDI X0, C.
+  MachineInstr &OffsetTail = *MRI->getVRegDef(Rs1);
+  if (OffsetTail.getOpcode() != RISCV::ADDI)
+    return false;
+  if (!OffsetTail.getOperand(1).isReg() ||
+      OffsetTail.getOperand(1).getReg() != RISCV::X0 ||
+      !OffsetTail.getOperand(2).isImm())
+    return false;
+
+  Offset = OffsetTail.getOperand(2).getImm();
+  switch (TailShXAdd.getOpcode()) {
+  default: llvm_unreachable("Unexpected opcode");
+  case RISCV::SH1ADD: Offset <<= 1; break;
+  case RISCV::SH2ADD: Offset <<= 2; break;
+  case RISCV::SH3ADD: Offset <<= 3; break;
+  }
+
+  LLVM_DEBUG(dbgs() << "  Offset Instr: " << OffsetTail);
+  DeadInstrs.insert(&OffsetTail);
+  return true;
+}
+
 bool RISCVMergeBaseOffsetOpt::detectAndFoldOffset(MachineInstr &HiLUI,
                                                   MachineInstr &LoADDI) {
   Register DestReg = LoADDI.getOperand(0).getReg();
   assert(MRI->hasOneUse(DestReg) && "expected one use for LoADDI");
   // LoADDI has only one use.
   MachineInstr &Tail = *MRI->use_instr_begin(DestReg);
   switch (Tail.getOpcode()) {
   default:
case RISCV::ADD: {
     // 2) LUI (offset20)
     //    This happens in case the lower 12 bits of the offset are zeros.
     int64_t Offset;
     if (!matchLargeOffset(Tail, DestReg, Offset))
       return false;
     foldOffset(HiLUI, LoADDI, Tail, Offset);
     return true;
   }
+  case RISCV::SH1ADD:
+  case RISCV::SH2ADD:
+  case RISCV::SH3ADD: {
+    // The offset is too large to fit in the immediate field of ADDI.
+    // It may be encoded as (SH2ADD (ADDI X0, C), DestReg) or
+    // (SH3ADD (ADDI X0, C), DestReg).
+    int64_t Offset;
+    if (!matchShiftedOffset(Tail, DestReg, Offset))
+      return false;
+    foldOffset(HiLUI, LoADDI, Tail, Offset);
+    return true;
+  }
   case RISCV::LB:
   case RISCV::LH:
   case RISCV::LW:
   case RISCV::LBU:
   case RISCV::LHU:
   case RISCV::LWU:
   case RISCV::LD:
   case RISCV::FLH:

llvm/test/CodeGen/RISCV/hoist-global-addr-base.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc -mtriple=riscv32  < %s | FileCheck  %s --check-prefixes=CHECK,RV32
-; RUN: llc -mtriple=riscv64  < %s | FileCheck  %s --check-prefixes=CHECK,RV64
+; RUN: llc -mtriple=riscv32  < %s | FileCheck  %s --check-prefixes=CHECK,RV32,RV32I
+; RUN: llc -mtriple=riscv64  < %s | FileCheck  %s --check-prefixes=CHECK,RV64,RV64I
+; RUN: llc -mtriple=riscv32 -mattr=+zba < %s | FileCheck  %s --check-prefixes=CHECK,RV32,RV32ZBA
+; RUN: llc -mtriple=riscv64 -mattr=+zba < %s | FileCheck  %s --check-prefixes=CHECK,RV64,RV64ZBA
 
 %struct.S = type { [40 x i32], i32, i32, i32, [4100 x i32], i32, i32, i32 }
 @s = common dso_local global %struct.S zeroinitializer, align 4
 @foo = global [6 x i16] [i16 1, i16 2, i16 3, i16 4, i16 5, i16 0], align 2
 @g = global [1048576 x i8] zeroinitializer, align 1
 @bar = external global [0 x i8], align 1
 
 

 define i8* @offset_addi_addi_neg() {
 ; CHECK-LABEL: offset_addi_addi_neg:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    lui a0, %hi(bar-4000)
 ; CHECK-NEXT:    addi a0, a0, %lo(bar-4000)
 ; CHECK-NEXT:    ret
     ret i8* getelementptr inbounds ([0 x i8], [0 x i8]* @bar, i32 0, i64 -4000)
 }
+
+; With Zba the constant 6424 is created with LI+SH2ADD.
+define i8* @offset_sh2add() {
+; CHECK-LABEL: offset_sh2add:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    lui a0, %hi(bar+6424)
+; CHECK-NEXT:    addi a0, a0, %lo(bar+6424)
+; CHECK-NEXT:    ret
+    ret i8* getelementptr inbounds ([0 x i8], [0 x i8]* @bar, i32 0, i64 6424)
+}
+
+; With Zba the constant 12848 is created with LI+SH3ADD.
+define i8* @offset_sh3add() {
+; CHECK-LABEL: offset_sh3add:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    lui a0, %hi(bar+12848)
+; CHECK-NEXT:    addi a0, a0, %lo(bar+12848)
+; CHECK-NEXT:    ret
+    ret i8* getelementptr inbounds ([0 x i8], [0 x i8]* @bar, i32 0, i64 12848)
+}
+;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
+; RV32I: {{.*}}
+; RV32ZBA: {{.*}}
+; RV64I: {{.*}}
+; RV64ZBA: {{.*}}