[llvm] [RISCV] Fix crash in PEI with empty entry block with Zcmp (PR #72117)

[Nemanja Ivanovic via llvm-commits]

https://github.com/nemanjai created https://github.com/llvm/llvm-project/pull/72117

We check the opcode of the first instruction in the block where the prologue is inserted without checking if the iterator points to any instructions. When the basic block is empty, that causes a crash. One way the prologue block can be empty is when it starts with a call to
__builtin_readcyclecounter on RV32 since that produces a loop.

>From 7424a9aa2fa57dc3dd222d1fd014bb3a61bf6eb4 Mon Sep 17 00:00:00 2001
From: Nemanja Ivanovic <nemanja at synopsys.com>
Date: Mon, 13 Nov 2023 15:06:00 +0100
Subject: [PATCH] [RISCV] Fix crash in PEI with empty entry block with Zcmp

We check the opcode of the first instruction in the block
where the prologue is inserted without checking if the
iterator points to any instructions. When the basic block
is empty, that causes a crash. One way the prologue block
can be empty is when it starts with a call to
__builtin_readcyclecounter on RV32 since that produces a loop.
---
 llvm/lib/Target/RISCV/RISCVFrameLowering.cpp |  3 ++-
 llvm/test/CodeGen/RISCV/pei-crash.ll         | 28 ++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/RISCV/pei-crash.ll

diff --git a/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp b/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp
index d873da7d684cfe2..1709ac28d3c5687 100644
--- a/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp
@@ -526,7 +526,8 @@ void RISCVFrameLowering::emitPrologue(MachineFunction &MF,
     RealStackSize = FirstSPAdjustAmount;
   }
 
-  if (RVFI->isPushable(MF) && FirstFrameSetup->getOpcode() == RISCV::CM_PUSH) {
+  if (RVFI->isPushable(MF) && FirstFrameSetup != MBB.end() &&
+      FirstFrameSetup->getOpcode() == RISCV::CM_PUSH) {
     // Use available stack adjustment in push instruction to allocate additional
     // stack space.
     uint64_t Spimm = std::min(StackSize, (uint64_t)48);
diff --git a/llvm/test/CodeGen/RISCV/pei-crash.ll b/llvm/test/CodeGen/RISCV/pei-crash.ll
new file mode 100644
index 000000000000000..7778f9580cf69df
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/pei-crash.ll
@@ -0,0 +1,28 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
+; RUN: llc -O0 --frame-pointer=none -mtriple=riscv32 -mattr=+zcmp \
+; RUN:   -verify-machineinstrs < %s | FileCheck %s
+define dso_local i64 @a() #0 {
+; CHECK-LABEL: a:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    addi sp, sp, -16
+; CHECK-NEXT:  .LBB0_1: # %entry
+; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    rdcycleh a0
+; CHECK-NEXT:    sw a0, 8(sp) # 4-byte Folded Spill
+; CHECK-NEXT:    rdcycle a1
+; CHECK-NEXT:    sw a1, 12(sp) # 4-byte Folded Spill
+; CHECK-NEXT:    rdcycleh a1
+; CHECK-NEXT:    bne a0, a1, .LBB0_1
+; CHECK-NEXT:  # %bb.2: # %entry
+; CHECK-NEXT:    lw a1, 8(sp) # 4-byte Folded Reload
+; CHECK-NEXT:    lw a0, 12(sp) # 4-byte Folded Reload
+; CHECK-NEXT:    addi sp, sp, 16
+; CHECK-NEXT:    ret
+entry:
+  %0 = call i64 @llvm.readcyclecounter()
+  ret i64 %0
+}
+
+declare i64 @llvm.readcyclecounter() #1
+
+attributes #0 = { noinline nounwind optnone }