[llvm] 53e89f5 - [RISCV] Add bounds check before use on returned iterator.

[Anmol P. Paralkar via llvm-commits]

Author: Anmol P. Paralkar
Date: 2023-08-14T16:06:09-07:00
New Revision: 53e89f5e3f2cefa432f7262343654c2a582dcbe1

URL: https://github.com/llvm/llvm-project/commit/53e89f5e3f2cefa432f7262343654c2a582dcbe1
DIFF: https://github.com/llvm/llvm-project/commit/53e89f5e3f2cefa432f7262343654c2a582dcbe1.diff

LOG: [RISCV] Add bounds check before use on returned iterator.

Check iterator validity before use; fixes a crash seen in the RISC-V
Zcmp Push/Pop optimization pass when compiling an internal benchmark.

Reviewed By: asb, wangpc

Differential Revision: https://reviews.llvm.org/D157674

Added: 
    llvm/test/CodeGen/RISCV/push-pop-opt-crash.ll

Modified: 
    llvm/lib/Target/RISCV/RISCVPushPopOptimizer.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/RISCV/RISCVPushPopOptimizer.cpp b/llvm/lib/Target/RISCV/RISCVPushPopOptimizer.cpp
index a93e750eadc6fa..f885adca669fe8 100644
--- a/llvm/lib/Target/RISCV/RISCVPushPopOptimizer.cpp
+++ b/llvm/lib/Target/RISCV/RISCVPushPopOptimizer.cpp
@@ -132,7 +132,8 @@ bool RISCVPushPopOpt::runOnMachineFunction(MachineFunction &Fn) {
   for (auto &MBB : Fn) {
     MachineBasicBlock::iterator MBBI = containsPop(MBB);
     MachineBasicBlock::iterator NextI = next_nodbg(MBBI, MBB.end());
-    if (MBBI != MBB.end() && NextI->getOpcode() == RISCV::PseudoRET)
+    if (MBBI != MBB.end() && NextI != MBB.end() &&
+        NextI->getOpcode() == RISCV::PseudoRET)
       Modified |= usePopRet(MBBI, NextI, adjustRetVal(MBBI));
   }
   return Modified;

diff  --git a/llvm/test/CodeGen/RISCV/push-pop-opt-crash.ll b/llvm/test/CodeGen/RISCV/push-pop-opt-crash.ll
new file mode 100644
index 00000000000000..ca6895dc51fa8c
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/push-pop-opt-crash.ll
@@ -0,0 +1,52 @@
+; RUN: llc -mattr=+zcmp -verify-machineinstrs  \
+; RUN: -mtriple=riscv32 -target-abi ilp32 < %s \
+; RUN: | FileCheck %s -check-prefixes=RV32IZCMP
+; RUN: llc -mattr=+zcmp -verify-machineinstrs  \
+; RUN: -mtriple=riscv64 -target-abi ilp64 < %s \
+; RUN: | FileCheck %s -check-prefixes=RV64IZCMP
+
+; This source code exposed a crash in the RISC-V Zcmp Push/Pop optimization
+; pass. The root cause was: Not doing a bounds check before using a returned
+; iterator.
+
+declare dso_local void @f1() local_unnamed_addr
+declare dso_local void @f2() local_unnamed_addr
+define  dso_local void @f0() local_unnamed_addr {
+; RV32IZCMP-LABEL: f0:
+; RV32IZCMP: 	.cfi_startproc
+; RV32IZCMP-NEXT: # %bb.0:                                # %entry
+; RV32IZCMP-NEXT: 	bnez	zero, .LBB0_2
+; RV32IZCMP-NEXT: # %bb.1:                                # %if.T
+; RV32IZCMP-NEXT: 	cm.push	{ra}, -16
+; RV32IZCMP-NEXT: 	.cfi_def_cfa_offset 16
+; RV32IZCMP-NEXT: 	.cfi_offset ra, -4
+; RV32IZCMP-NEXT: 	call	f1
+; RV32IZCMP-NEXT: 	cm.pop	{ra}, 16
+; RV32IZCMP-NEXT: .LBB0_2:                                # %if.F
+; RV32IZCMP-NEXT: 	tail	f2
+; RV32IZCMP-NEXT: .Lfunc_end0:
+
+; RV64IZCMP-LABEL: f0:
+; RV64IZCMP: 	.cfi_startproc
+; RV64IZCMP-NEXT: # %bb.0:                                # %entry
+; RV64IZCMP-NEXT: 	bnez	zero, .LBB0_2
+; RV64IZCMP-NEXT: # %bb.1:                                # %if.T
+; RV64IZCMP-NEXT: 	cm.push	{ra}, -16
+; RV64IZCMP-NEXT: 	.cfi_def_cfa_offset 16
+; RV64IZCMP-NEXT: 	.cfi_offset ra, -8
+; RV64IZCMP-NEXT: 	call	f1
+; RV64IZCMP-NEXT: 	cm.pop	{ra}, 16
+; RV64IZCMP-NEXT: .LBB0_2:                                # %if.F
+; RV64IZCMP-NEXT: 	tail	f2
+; RV64IZCMP-NEXT: .Lfunc_end0:
+entry:
+  br i1 poison, label %if.T, label %if.F
+
+if.T:
+  tail call void @f1()
+  br label %if.F
+
+if.F:
+  tail call void @f2()
+  ret void
+}