From 6564e0cf7e61518cb15443fca42bc2206a6123e2 Mon Sep 17 00:00:00 2001
From: Yonghong Song <yhs@fb.com>
Date: Tue, 2 Mar 2021 09:35:21 -0800
Subject: [PATCH] BPF: Fix a bug in peephole TRUNC elimination optimization

Andrei Matei reported a llvm11 core dump for his bpf program
   https://bugs.llvm.org/show_bug.cgi?id=48578
The core dump happens in LiveVariables analysis phase.
  #4 0x00007fce54356bb0 __restore_rt
  #5 0x00007fce4d51785e llvm::LiveVariables::HandleVirtRegUse(unsigned int,
      llvm::MachineBasicBlock*, llvm::MachineInstr&)
  #6 0x00007fce4d519abe llvm::LiveVariables::runOnInstr(llvm::MachineInstr&,
      llvm::SmallVectorImpl<unsigned int>&)
  #7 0x00007fce4d519ec6 llvm::LiveVariables::runOnBlock(llvm::MachineBasicBlock*, unsigned int)
  #8 0x00007fce4d51a4bf llvm::LiveVariables::runOnMachineFunction(llvm::MachineFunction&)
The bug can be reproduced with llvm12 and latest trunk as well.

Futher analysis shows that there is a bug in BPF peephole
TRUNC elimination optimization, which tries to remove
unnecessary TRUNC operations (a <<= 32; a >>= 32).
Specifically, the compiler did wrong transformation for the
following patterns:
   %1 = LDW ...
   %2 = SLL_ri %1, 32
   %3 = SRL_ri %2, 32
   ... %3 ...
   %4 = SRA_ri %2, 32
   ... %4 ...

The current transformation did not check how many uses of %2
and did transformation like
   %1 = LDW ...
   ... %1 ...
   %4 = SRL_ri %2, 32
   ... %4 ...
and pseudo register %2 is used by not defined and
caused LiveVariables analysis core dump.

To fix the issue, when traversing back from SRL_ri to SLL_ri,
check to ensure SLL_ri has only one use. Otherwise, don't
do transformation.

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

(cherry picked from commit 51cdb780db3b9b46c783efcec672c4da272e9992)
---
 llvm/lib/Target/BPF/BPFMIPeephole.cpp      |  3 ++
 llvm/test/CodeGen/BPF/remove_truncate_8.ll | 41 ++++++++++++++++++++++
 2 files changed, 44 insertions(+)
 create mode 100644 llvm/test/CodeGen/BPF/remove_truncate_8.ll

diff --git a/llvm/lib/Target/BPF/BPFMIPeephole.cpp b/llvm/lib/Target/BPF/BPFMIPeephole.cpp
index df870314fffe3b..354980e4bf3ce1 100644
--- a/llvm/lib/Target/BPF/BPFMIPeephole.cpp
+++ b/llvm/lib/Target/BPF/BPFMIPeephole.cpp
@@ -475,6 +475,9 @@ bool BPFMIPeepholeTruncElim::eliminateTruncSeq(void) {
       if (MI.getOpcode() == BPF::SRL_ri &&
           MI.getOperand(2).getImm() == 32) {
         SrcReg = MI.getOperand(1).getReg();
+        if (!MRI->hasOneNonDBGUse(SrcReg))
+          continue;
+
         MI2 = MRI->getVRegDef(SrcReg);
         DstReg = MI.getOperand(0).getReg();
 
diff --git a/llvm/test/CodeGen/BPF/remove_truncate_8.ll b/llvm/test/CodeGen/BPF/remove_truncate_8.ll
new file mode 100644
index 00000000000000..fb1eabb0f0fd8d
--- /dev/null
+++ b/llvm/test/CodeGen/BPF/remove_truncate_8.ll
@@ -0,0 +1,41 @@
+; RUN: llc < %s -march=bpf -verify-machineinstrs | FileCheck %s
+; Source Code:
+;   struct loc_prog {
+;     unsigned int ip;
+;     int len;
+;   };
+;   int exec_prog(struct loc_prog *prog) {
+;     if (prog->ip < prog->len) {
+;       int x = prog->ip;
+;       if (x < 3)
+;         prog->ip += 2;
+;     }
+;     return 3;
+;   }
+; Compilation flag:
+;   clang -target bpf -O2 -S -emit-llvm t.c
+
+%struct.loc_prog = type { i32, i32 }
+
+; Function Attrs: nofree norecurse nounwind willreturn
+define dso_local i32 @exec_prog(%struct.loc_prog* nocapture %prog) local_unnamed_addr {
+entry:
+  %ip = getelementptr inbounds %struct.loc_prog, %struct.loc_prog* %prog, i64 0, i32 0
+  %0 = load i32, i32* %ip, align 4
+  %len = getelementptr inbounds %struct.loc_prog, %struct.loc_prog* %prog, i64 0, i32 1
+  %1 = load i32, i32* %len, align 4
+  %cmp = icmp ult i32 %0, %1
+  %cmp2 = icmp slt i32 %0, 3
+  %or.cond = and i1 %cmp2, %cmp
+; CHECK: r{{[0-9]+}} <<= 32
+; CHECK: r{{[0-9]+}} s>>= 32
+  br i1 %or.cond, label %if.then3, label %if.end5
+
+if.then3:                                         ; preds = %entry
+  %add = add nsw i32 %0, 2
+  store i32 %add, i32* %ip, align 4
+  br label %if.end5
+
+if.end5:                                          ; preds = %if.then3, %entry
+  ret i32 3
+}