BPF: avoid same condition evaluated in multiple branch instructions

[yonghong-song]

Summary

This patch specifically tries to address an issue reported at

https://lore.kernel.org/bpf/20200812074826.GB754656@krava/

If one icmp condition is used in more than one places, prevent
this condition from inlining into multiple places.

For example,

cond = a < b 
... 
select cond, val1, val2
... 
br cond, block1, block2

In actually lowering, cond will be put into select_cc and br_cc
so we will have proper cond in the jmp inst.
It is reported this caused the issue if cond happens to be

pkt_ptr < pkt_end

due to particular kernel implementation.
The transformation here is to put a barrier after "cond = a < b"
so cond cannot be inlined into select/br. The pseudo generated code
looks like:

if a < b then
  tmp = 1 
then
  tmp = 0 
... 
if (tmp == 1) val1; else val2
... 
if (tmp == 1) goto block1; else goto block2;

The below llc flag is added to turn off this optimization if needed.

llc -march=bpf -bpf-disable-eval-cond-once t.ll

[yonghong-song]

rebase since https://reviews.llvm.org/D87153 changed

[yonghong-song]

• added tests, added llc option to disable optimization if needed.
• removed RFC tag.

llvm/lib/Target/BPF/BPFCheckAndAdjustIR.cpp

 #include "llvm/IR/Value.h"
 #include "llvm/Pass.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 
 #define DEBUG_TYPE "bpf-check-and-opt-ir"
 
 using namespace llvm;
 
+static cl::opt<bool> DisableBPFevalCondOnce(
+    "bpf-disable-eval-cond-once", cl::Hidden,
+    cl::desc("BPF: Disable Evaluating Conditionals Once"),
+    cl::init(false));
+
 namespace {
 
 class BPFCheckAndAdjustIR final : public ModulePass {
   bool runOnModule(Module &F) override;
 
 public:
   static char ID;
   BPFCheckAndAdjustIR() : ModulePass(ID) {}
 
 private:
   void checkIR(Module &M);
   bool adjustIR(Module &M);
   bool removePassThroughBuiltin(Module &M);
+  bool evalCondOnce(Module &M);
 };
 } // End anonymous namespace
 
 char BPFCheckAndAdjustIR::ID = 0;
 INITIALIZE_PASS(BPFCheckAndAdjustIR, DEBUG_TYPE, "BPF Check And Adjust IR",
                 false, false)
 
 ModulePass *llvm::createBPFCheckAndAdjustIR() {
for (auto &BB : F)
         Changed = true;
         Value *Arg = Call->getArgOperand(1);
         Call->replaceAllUsesWith(Arg);
         ToBeDeleted = Call;
       }
   return Changed;
 }
 
+bool BPFCheckAndAdjustIR::evalCondOnce(Module &M) {
+  // Forcing to lower an icmp condition if it is used
+  // in multiple places.
+  // For:
+  //   icmp1 = ...
+  //   ...
+  //   a = select icmp1, val1, val2
+  //   ...
+  //   ... icmp1 ...
+  // Change to:
+  //   icmp1 = ...
+  //   icmp1 = user_asm_barrier(icmp1)
+  //   ...
+  //   a = select icmp1, val1, val2
+  //   ...
+  //   ... icmp1 ...
+  SmallVector<ICmpInst *, 16> Candidates;
+  for (Function &F : M)
+    for (auto &BB : F)
+      for (auto &I : BB) {
+        auto *Cond = dyn_cast<ICmpInst>(&I);
+        if (!Cond || Cond->use_empty() || Cond->hasOneUse())
+          continue;
+        Candidates.push_back(Cond);
+      }
+
+  if (Candidates.empty())
+    return false;
+
+  bool Changed = false;
+  for (auto *Cond : Candidates) {
+    SmallVector<Instruction *, 16> Uses;
+    for (User *U : Cond->users()) {
+      Instruction *Inst = dyn_cast<Instruction>(U);
+      if (!Inst || (!isa<SelectInst>(Inst) && !isa<BranchInst>(Inst)))
+        continue;
+      Uses.push_back(Inst);
+    }
+    if (Uses.size() < 2)
+      continue;
+
+    Changed = true;
+    Instruction *NextI = Cond->getNextNonDebugInstruction();
+    assert(NextI);
+    FunctionType *AsmTy = FunctionType::get(Cond->getType(), {Cond->getType()}, false);
+    InlineAsm *Asm = InlineAsm::get(AsmTy, StringRef(""), StringRef("=r,0"), true);
+    auto *CI = CallInst::Create(Asm, {Cond}, "", NextI);
+    for (Instruction *Inst : Uses) {
+      Inst->setOperand(0, CI);
+    }
+  }
+
+  return Changed;
+}
+
 bool BPFCheckAndAdjustIR::adjustIR(Module &M) {
-  return removePassThroughBuiltin(M);
+  bool Changed = false;
+
+  Changed = removePassThroughBuiltin(M);
+  if (!DisableBPFevalCondOnce)
+    Changed = evalCondOnce(M) || Changed;
+  return Changed;
 }
 
 bool BPFCheckAndAdjustIR::runOnModule(Module &M) {
   checkIR(M);
   return adjustIR(M);
 }

llvm/test/CodeGen/BPF/eval-cond-once-1.ll

+; RUN: llc < %s -march=bpfel | FileCheck %s
+;
+; Source:
+;   struct __sk_buff { void *data; void *data_end; };
+;   struct ethhdr { char unused[14]; };
+;   struct iphdr { char unused[40]; short protocol; };
+;   struct udphdr { char unused[20]; unsigned dest; };
+;
+;   static __attribute__((always_inline)) struct iphdr *get_iphdr(struct __sk_buff *skb)
+;   {
+;     struct iphdr *ip = 0;
+;     struct ethhdr *eth;
+;
+;     if (skb->data + 42 > skb->data_end)
+;       goto out;
+;
+;     eth = (void *)skb->data;
+;     ip = (void *)(eth + 14);
+;
+;   out:
+;     return ip;
+;  }
+;
+;  int test(struct __sk_buff *skb) {
+;    struct iphdr *ip = 0;
+;    struct udphdr *udp;
+;    short proto = 0;
+;
+;    if (!(ip = get_iphdr(skb)))
+;      goto out;
+;
+;    proto = ip->protocol;
+;    if (proto != 4)
+;      goto out;
+;
+;    udp = (void *)(ip + 1);
+;    if (udp->dest != 10)
+;      goto out;
+;
+;    return udp->dest;
+;
+;  out:
+;    return -1;
+;  }
+; Compilation flag:
+;   clang -target bpf -O2 -S -emit-llvm t.c
+
+%struct.__sk_buff = type { i8*, i8* }
+%struct.iphdr = type { [40 x i8], i16 }
+
+; Function Attrs: norecurse nounwind readonly
+define dso_local i32 @test(%struct.__sk_buff* nocapture readonly %skb) local_unnamed_addr {
+entry:
+  %data.i = getelementptr inbounds %struct.__sk_buff, %struct.__sk_buff* %skb, i64 0, i32 0
+  %0 = load i8*, i8** %data.i, align 8
+  %add.ptr.i = getelementptr i8, i8* %0, i64 42
+  %data_end.i = getelementptr inbounds %struct.__sk_buff, %struct.__sk_buff* %skb, i64 0, i32 1
+  %1 = load i8*, i8** %data_end.i, align 8
+  %cmp.i = icmp ugt i8* %add.ptr.i, %1
+  %add.ptr2.i = getelementptr inbounds i8, i8* %0, i64 196
+  %2 = bitcast i8* %add.ptr2.i to %struct.iphdr*
+  %ip.0.i = select i1 %cmp.i, %struct.iphdr* null, %struct.iphdr* %2
+  br i1 %cmp.i, label %out, label %if.end
+
+; CHECK:        [[REG3:r[0-9]+]] = 1
+; CHECK-NEXT:   if r{{[0-9]+}} > r{{[0-9]+}} goto [[LABEL:.*]]
+; CHECK:        [[REG3]] = 0
+; CHECK:        [[LABEL]]:
+; CHECK:        if [[REG3]] != 0 goto
+
+if.end:                                           ; preds = %entry
+  %protocol = getelementptr inbounds %struct.iphdr, %struct.iphdr* %ip.0.i, i64 0, i32 1
+  %3 = load i16, i16* %protocol, align 2
+  %cmp.not = icmp eq i16 %3, 4
+  br i1 %cmp.not, label %if.end3, label %out
+
+; CHECK:        if [[REG3]] != 0 goto
+
+if.end3:                                          ; preds = %if.end
+  %dest = getelementptr inbounds %struct.iphdr, %struct.iphdr* %ip.0.i, i64 1, i32 0, i64 20
+  %4 = bitcast i8* %dest to i32*
+  %5 = load i32, i32* %4, align 4
+  %cmp4.not = icmp eq i32 %5, 10
+  br i1 %cmp4.not, label %cleanup, label %out
+
+out:                                              ; preds = %if.end3, %if.end, %entry
+  br label %cleanup
+
+cleanup:                                          ; preds = %if.end3, %out
+  %retval.0 = phi i32 [ -1, %out ], [ 10, %if.end3 ]
+  ret i32 %retval.0
+}

llvm/test/CodeGen/BPF/eval-cond-once-2.ll

+; RUN: llc < %s -march=bpfel | FileCheck %s
+;
+; Source:
+;   struct s1 { void *data; void *data_end; };
+;   void foo(int);
+;   int test(struct s1 *arg) {
+;     int i, cond = arg->data + 42 < arg->data_end;
+;     if (cond) foo(-1);
+;     for (i = 0; i < 100; i++)
+;       foo(i);
+;     if (cond) foo(100);
+;     return 0;
+;   }
+; Compilation flag:
+;    clang -target bpf -O2 -S -emit-llvm t.c
+
+%struct.s1 = type { i8*, i8* }
+
+; Function Attrs: nounwind
+define dso_local i32 @test(%struct.s1* nocapture readonly %arg) local_unnamed_addr {
+entry:
+  %data = getelementptr inbounds %struct.s1, %struct.s1* %arg, i64 0, i32 0
+  %0 = load i8*, i8** %data, align 8
+  %add.ptr = getelementptr i8, i8* %0, i64 42
+  %data_end = getelementptr inbounds %struct.s1, %struct.s1* %arg, i64 0, i32 1
+  %1 = load i8*, i8** %data_end, align 8
+  %cmp = icmp ult i8* %add.ptr, %1
+  br i1 %cmp, label %if.then, label %for.body.preheader
+
+; CHECK:        [[REG7:r[0-9]+]] = 1
+; CHECK-NEXT:   if r{{[0-9]+}} > r{{[0-9]+}} goto [[LABEL:.*]]
+; CHECK:        [[REG7]] = 0
+; CHECK:        [[LABEL]]:
+
+; CHECK:        [[REG1:r[0-9]+]] = [[REG7]]
+; CHECK:        [[REG1]] &= 1
+; CHECK-NEXT:   if [[REG1]] == 0 goto
+
+if.then:                                          ; preds = %entry
+  tail call void @foo(i32 -1)
+  br label %for.body.preheader
+
+for.body.preheader:                               ; preds = %if.then, %entry
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.preheader, %for.body
+  %i.012 = phi i32 [ %inc, %for.body ], [ 0, %for.body.preheader ]
+  tail call void @foo(i32 %i.012)
+  %inc = add nuw nsw i32 %i.012, 1
+  %exitcond.not = icmp eq i32 %inc, 100
+  br i1 %exitcond.not, label %for.end, label %for.body
+
+for.end:                                          ; preds = %for.body
+  br i1 %cmp, label %if.then4, label %if.end5
+
+; CHECK:        [[REG7]] &= 1
+; CHECK-NEXT:   if [[REG7]] == 0 goto
+
+if.then4:                                         ; preds = %for.end
+  tail call void @foo(i32 100)
+  br label %if.end5
+
+if.end5:                                          ; preds = %if.then4, %for.end
+  ret i32 0
+}
+
+declare dso_local void @foo(i32) local_unnamed_addr