BPF: Workaround an InstCombine ICmp transformation with llvm.bpf.compare builtin

[yonghong-song]

Summary

Commit acabad9ff6bf ("[InstCombine] try to canonicalize icmp with
trunc op into mask and cmp") added a transformation to
convert "(conv)a < power_2_const" to "a & <const>" in certain
cases and bpf kernel verifier has to handle the resulted code
conservatively and this may reject otherwise legitimate program.

This commit tries to prevent such a transformation. A bpf backend
builtin llvm.bpf.compare is added. The ICMP insn, which is subject to
above InstCombine transformation, is converted to the builtin
function. The builtin function is later lowered to original ICMP insn,
certainly after InstCombine pass.

With this change, all affected bpf strobemeta* selftests are
passed now.

llvm/include/llvm/IR/IntrinsicsBPF.td

let TargetPrefix = "bpf" in {  // All intrinsics start with "llvm.bpf."
   def int_bpf_preserve_type_info : GCCBuiltin<"__builtin_bpf_preserve_type_info">,
               Intrinsic<[llvm_i32_ty], [llvm_i32_ty, llvm_i64_ty],
               [IntrNoMem]>;
   def int_bpf_preserve_enum_value : GCCBuiltin<"__builtin_bpf_preserve_enum_value">,
               Intrinsic<[llvm_i64_ty], [llvm_i32_ty, llvm_ptr_ty, llvm_i64_ty],
               [IntrNoMem]>;
   def int_bpf_passthrough : GCCBuiltin<"__builtin_bpf_passthrough">,
               Intrinsic<[llvm_any_ty], [llvm_i32_ty, llvm_any_ty], [IntrNoMem]>;
+  def int_bpf_compare : GCCBuiltin<"__builtin_bpf_compare">,
+              Intrinsic<[llvm_i1_ty], [llvm_i32_ty, llvm_anyint_ty, llvm_anyint_ty],
+              [IntrNoMem]>;
 }

llvm/lib/Target/BPF/BPFAdjustOpt.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "BPF.h"
 #include "BPFCORE.h"
 #include "BPFTargetMachine.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicsBPF.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Pass.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 
public:
   BPFAdjustOptImpl(Module *M) : M(M) {}
 
   bool run();
 
 private:
   Module *M;
   SmallVector<PassThroughInfo, 16> PassThroughs;
 
+  bool adjustICmpToBuiltin();
   void adjustBasicBlock(BasicBlock &BB);
   bool serializeICMPCrossBB(BasicBlock &BB);
   void adjustInst(Instruction &I);
   bool serializeICMPInBB(Instruction &I);
   bool avoidSpeculation(Instruction &I);
   bool insertPassThrough();
 };
 
 } // End anonymous namespace
 
 char BPFAdjustOpt::ID = 0;
 INITIALIZE_PASS(BPFAdjustOpt, "bpf-adjust-opt", "BPF Adjust Optimization",
                 false, false)
 
 ModulePass *llvm::createBPFAdjustOpt() { return new BPFAdjustOpt(); }
 
 bool BPFAdjustOpt::runOnModule(Module &M) { return BPFAdjustOptImpl(&M).run(); }
 
 bool BPFAdjustOptImpl::run() {
+  bool Changed = adjustICmpToBuiltin();
+
   for (Function &F : *M)
     for (auto &BB : F) {
       adjustBasicBlock(BB);
       for (auto &I : BB)
         adjustInst(I);
     }
+  return insertPassThrough() || Changed;
+}
+
+// Commit acabad9ff6bf ("[InstCombine] try to canonicalize icmp with
+// trunc op into mask and cmp") added a transformation to
+// convert "(conv)a < power_2_const" to "a & <const>" in certain
+// cases and bpf kernel verifier has to handle the resulted code
+// conservatively and this may reject otherwise legitimate program.
+// Here, we change related icmp code to a builtin which will
+// be restored to original icmp code later to prevent that
+// InstCombine transformatin.
+bool BPFAdjustOptImpl::adjustICmpToBuiltin() {
+  bool Changed = false;
+  ICmpInst *ToBeDeleted = nullptr;
+  for (Function &F : *M)
+    for (auto &BB : F)
+      for (auto &I : BB) {
+        if (ToBeDeleted) {
+          ToBeDeleted->eraseFromParent();
+          ToBeDeleted = nullptr;
+        }
+
+        auto *Icmp = dyn_cast<ICmpInst>(&I);
+        if (!Icmp)
+          continue;
+
+        Value *Op0 = Icmp->getOperand(0);
+        if (!isa<TruncInst>(Op0))
+          continue;
+
+        auto ConstOp1 = dyn_cast<ConstantInt>(Icmp->getOperand(1));
+        if (!ConstOp1)
+          continue;
+
+        auto ConstOp1Val = ConstOp1->getValue().getZExtValue();
+        auto Op = Icmp->getPredicate();
+        if (Op == ICmpInst::ICMP_ULT) {
+          if ((ConstOp1Val - 1) & ConstOp1Val)
+            continue;
+        } else if (Op == ICmpInst::ICMP_ULE) {
+          if (ConstOp1Val & (ConstOp1Val + 1))
+            continue;
+        } else {
+          continue;
+        }
+
+        Constant *Opcode =
+            ConstantInt::get(Type::getInt32Ty(BB.getContext()), Op);
+        Function *Fn = Intrinsic::getDeclaration(
+            M, Intrinsic::bpf_compare, {Op0->getType(), ConstOp1->getType()});
+        auto *NewInst = CallInst::Create(Fn, {Opcode, Op0, ConstOp1});
+        BB.getInstList().insert(I.getIterator(), NewInst);
+        Icmp->replaceAllUsesWith(NewInst);
+        Changed = true;
+        ToBeDeleted = Icmp;
+      }
 
-  return insertPassThrough();
+  return Changed;
 }
 
 bool BPFAdjustOptImpl::insertPassThrough() {
   for (auto &Info : PassThroughs) {
     auto *CI = BPFCoreSharedInfo::insertPassThrough(
         M, Info.UsedInst->getParent(), Info.Input, Info.UsedInst);
     Info.UsedInst->setOperand(Info.OpIdx, CI);
   }

llvm/lib/Target/BPF/BPFCheckAndAdjustIR.cpp

 public:
   static char ID;
   BPFCheckAndAdjustIR() : ModulePass(ID) {}
 
 private:
   void checkIR(Module &M);
   bool adjustIR(Module &M);
   bool removePassThroughBuiltin(Module &M);
+  bool removeCompareBuiltin(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::removeCompareBuiltin(Module &M) {
+  // Remove __builtin_bpf_compare()'s which are used to prevent
+  // certain IR optimizations. Now major IR optimizations are done,
+  // remove them.
+  bool Changed = false;
+  CallInst *ToBeDeleted = nullptr;
+  for (Function &F : M)
+    for (auto &BB : F)
+      for (auto &I : BB) {
+        if (ToBeDeleted) {
+          ToBeDeleted->eraseFromParent();
+          ToBeDeleted = nullptr;
+        }
+
+        auto *Call = dyn_cast<CallInst>(&I);
+        if (!Call)
+          continue;
+        auto *GV = dyn_cast<GlobalValue>(Call->getCalledOperand());
+        if (!GV)
+          continue;
+        if (!GV->getName().startswith("llvm.bpf.compare"))
+          continue;
+
+        Changed = true;
+        Value *Arg0 = Call->getArgOperand(0);
+        Value *Arg1 = Call->getArgOperand(1);
+        Value *Arg2 = Call->getArgOperand(2);
+
+        auto OpVal = cast<ConstantInt>(Arg0)->getValue().getZExtValue();
+        CmpInst::Predicate Opcode = (CmpInst::Predicate)OpVal;
+
+        auto *ICmp = new ICmpInst(Opcode, Arg1, Arg2);
+        BB.getInstList().insert(Call->getIterator(), ICmp);
+
+        Call->replaceAllUsesWith(ICmp);
+        ToBeDeleted = Call;
+      }
+  return Changed;
+}
+
 bool BPFCheckAndAdjustIR::adjustIR(Module &M) {
-  return removePassThroughBuiltin(M);
+  bool Changed = removePassThroughBuiltin(M);
+  Changed = removeCompareBuiltin(M) || Changed;
+  return Changed;
 }
 
 bool BPFCheckAndAdjustIR::runOnModule(Module &M) {
   checkIR(M);
   return adjustIR(M);
 }

llvm/test/CodeGen/BPF/adjust-opt-icmp3.ll

+; RUN: opt -O2 -S -mtriple=bpf-pc-linux %s -o %t1
+; RUN: llc %t1 -o - | FileCheck -check-prefixes=CHECK,CHECK-V1 %s
+; RUN: opt -O2 -S -mtriple=bpf-pc-linux %s -o %t1
+; RUN: llc %t1 -mcpu=v3 -o - | FileCheck -check-prefixes=CHECK,CHECK-V3 %s
+;
+; Source:
+;   int test1(unsigned long a) {
+;     if ((unsigned)a <= 3) return 2;
+;     return 3;
+;   }
+;   int test2(unsigned long a) {
+;     if ((unsigned)a < 4) return 2;
+;     return 3;
+;   }
+; Compilation flag:
+;   clang -target bpf -O2 -S -emit-llvm -Xclang -disable-llvm-passes test.c
+
+; Function Attrs: nounwind
+define dso_local i32 @test1(i64 %a) #0 {
+entry:
+  %retval = alloca i32, align 4
+  %a.addr = alloca i64, align 8
+  store i64 %a, i64* %a.addr, align 8, !tbaa !3
+  %0 = load i64, i64* %a.addr, align 8, !tbaa !3
+  %conv = trunc i64 %0 to i32
+  %cmp = icmp ule i32 %conv, 3
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  store i32 2, i32* %retval, align 4
+  br label %return
+
+if.end:                                           ; preds = %entry
+  store i32 3, i32* %retval, align 4
+  br label %return
+
+return:                                           ; preds = %if.end, %if.then
+  %1 = load i32, i32* %retval, align 4
+  ret i32 %1
+}
+
+; CHECK-LABEL: test1
+; CHECK-V1:    if r[[#]] > r[[#]] goto
+; CHECK-V3:    if w[[#]] < 4 goto
+
+; Function Attrs: nounwind
+define dso_local i32 @test2(i64 %a) #0 {
+entry:
+  %retval = alloca i32, align 4
+  %a.addr = alloca i64, align 8
+  store i64 %a, i64* %a.addr, align 8, !tbaa !3
+  %0 = load i64, i64* %a.addr, align 8, !tbaa !3
+  %conv = trunc i64 %0 to i32
+  %cmp = icmp ult i32 %conv, 4
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  store i32 2, i32* %retval, align 4
+  br label %return
+
+if.end:                                           ; preds = %entry
+  store i32 3, i32* %retval, align 4
+  br label %return
+
+return:                                           ; preds = %if.end, %if.then
+  %1 = load i32, i32* %retval, align 4
+  ret i32 %1
+}
+
+; CHECK-LABEL: test2
+; CHECK-V1:    if r[[#]] > r[[#]] goto
+; CHECK-V3:    if w[[#]] < 4 goto
+
+attributes #0 = { nounwind "frame-pointer"="all" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" }
+
+!llvm.module.flags = !{!0, !1}
+!llvm.ident = !{!2}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 7, !"frame-pointer", i32 2}
+!2 = !{!"clang version 14.0.0 (https://github.com/llvm/llvm-project.git b7892f95881c891032742e0cd81861b845512653)"}
+!3 = !{!4, !4, i64 0}
+!4 = !{!"long", !5, i64 0}
+!5 = !{!"omnipotent char", !6, i64 0}
+!6 = !{!"Simple C/C++ TBAA"}