[llvm] [RISCV] Fix assertion failure in `genShXAddAddShift` (PR #88757)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/88757

Fix assertion failure in our downstream CI https://github.com/dtcxzyw/llvm-codegen-benchmark/issues/1.


>From 44f7ace1a5852c100a42bd90e3269bc969a34106 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Tue, 16 Apr 2024 00:46:48 +0800
Subject: [PATCH] [RISCV] Fix assertion failure in `genShXAddAddShift`

---
 llvm/lib/Target/RISCV/RISCVInstrInfo.cpp |  2 +-
 llvm/test/CodeGen/RISCV/rv64zba.ll       | 23 +++++++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp b/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp
index 508f607fab20fd..b0fda040519a57 100644
--- a/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp
+++ b/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp
@@ -1986,7 +1986,7 @@ genShXAddAddShift(MachineInstr &Root, unsigned AddOpIdx,
       MRI.getUniqueVRegDef(AddMI->getOperand(AddOpIdx).getReg());
 
   unsigned InnerShiftAmt = ShiftMI->getOperand(2).getImm();
-  assert(InnerShiftAmt > OuterShiftAmt && "Unexpected shift amount");
+  assert(InnerShiftAmt >= OuterShiftAmt && "Unexpected shift amount");
 
   unsigned InnerOpc;
   switch (InnerShiftAmt - OuterShiftAmt) {
diff --git a/llvm/test/CodeGen/RISCV/rv64zba.ll b/llvm/test/CodeGen/RISCV/rv64zba.ll
index 067addc819f7e6..bb4be323ecb2e5 100644
--- a/llvm/test/CodeGen/RISCV/rv64zba.ll
+++ b/llvm/test/CodeGen/RISCV/rv64zba.ll
@@ -2389,3 +2389,26 @@ define i64 @array_index_sh4_sh3(ptr %p, i64 %idx1, i64 %idx2) {
   %b = load i64, ptr %a, align 8
   ret i64 %b
 }
+
+define ptr @test_gep_gep_dont_crash(ptr %p, i64 %a1, i64 %a2) {
+; RV64I-LABEL: test_gep_gep_dont_crash:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    srliw a2, a2, 6
+; RV64I-NEXT:    slli a2, a2, 3
+; RV64I-NEXT:    slli a1, a1, 3
+; RV64I-NEXT:    add a0, a0, a1
+; RV64I-NEXT:    add a0, a0, a2
+; RV64I-NEXT:    ret
+;
+; RV64ZBA-LABEL: test_gep_gep_dont_crash:
+; RV64ZBA:       # %bb.0:
+; RV64ZBA-NEXT:    srliw a2, a2, 6
+; RV64ZBA-NEXT:    add a1, a2, a1
+; RV64ZBA-NEXT:    sh3add a0, a1, a0
+; RV64ZBA-NEXT:    ret
+  %lshr = lshr i64 %a2, 6
+  %and = and i64 %lshr, 67108863
+  %gep1 = getelementptr i64, ptr %p, i64 %and
+  %gep2 = getelementptr i64, ptr %gep1, i64 %a1
+  ret ptr %gep2
+}