[Clang][BPF]: Force sign/zero extension for return values in caller

Currently bpf supports calling kernel functions (x86_64, arm64, etc.)
in bpf programs. Tejun discovered a problem where the x86_64 func
return value (a unsigned char type) is stored in 8-bit subregister %al
and the other 56-bits in %rax might be garbage. But based on current
bpf ABI, the bpf program assumes the whole %rax holds the correct value
as the callee is supposed to do necessary sign/zero extension.
This mismatch between bpf and x86_64 caused the incorrect results.

To resolve this problem, this patch forced caller to do needed
sign/zero extension for 8/16-bit return values as well. Note that
32-bit return values already had sign/zero extension even without
this patch.

For example, for the test case attached to this patch:

  $  cat t.c
  _Bool bar_bool(void);
  unsigned char bar_char(void);
  short bar_short(void);
  int bar_int(void);
  int foo_bool(void) {
        if (bar_bool() != 1) return 0; else return 1;
  }
  int foo_char(void) {
        if (bar_char() != 10) return 0; else return 1;
  }
  int foo_short(void) {
        if (bar_short() != 10) return 0; else return 1;
  }
  int foo_int(void) {
        if (bar_int() != 10) return 0; else return 1;
  }

Without this patch, generated call insns in IR looks like:
    %call = call zeroext i1 @bar_bool()
    %call = call zeroext i8 @bar_char()
    %call = call signext i16 @bar_short()
    %call = call i32 @bar_int()
So it is assumed that zero extension has been done for return values of
bar_bool()and bar_char(). Sign extension has been done for the return
value of bar_short(). The return value of bar_int() does not have any
assumption so caller needs to do necessary shifting to get correct
32bit values.

With this patch, generated call insns in IR looks like:
    %call = call i1 @bar_bool()
    %call = call i8 @bar_char()
    %call = call i16 @bar_short()
    %call = call i32 @bar_int()
There are no assumptions for return values of the above four function calls,
so necessary shifting is necessary for all of them.

The following is the objdump file difference for function foo_char().
Without this patch:
  0000000000000010 <foo_char>:
       2:       85 10 00 00 ff ff ff ff call -1
       3:       bf 01 00 00 00 00 00 00 r1 = r0
       4:       b7 00 00 00 01 00 00 00 r0 = 1
       5:       15 01 01 00 0a 00 00 00 if r1 == 10 goto +1 <LBB1_2>
       6:       b7 00 00 00 00 00 00 00 r0 = 0
  0000000000000038 <LBB1_2>:
       7:       95 00 00 00 00 00 00 00 exit

With this patch:
  0000000000000018 <foo_char>:
       3:       85 10 00 00 ff ff ff ff call -1
       4:       bf 01 00 00 00 00 00 00 r1 = r0
       5:       57 01 00 00 ff 00 00 00 r1 &= 255
       6:       b7 00 00 00 01 00 00 00 r0 = 1
       7:       15 01 01 00 0a 00 00 00 if r1 == 10 goto +1 <LBB1_2>
       8:       b7 00 00 00 00 00 00 00 r0 = 0
  0000000000000048 <LBB1_2>:
       9:       95 00 00 00 00 00 00 00 exit
The zero extension of the return 'char' value is done here.

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

clang/lib/CodeGen/TargetInfo.cpp

@@ -11499,6 +11499,56 @@ class CSKYTargetCodeGenInfo : public TargetCodeGenInfo {
 };
 } // end anonymous namespace
 
+//===----------------------------------------------------------------------===//
+// BPF ABI Implementation
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+class BPFABIInfo : public DefaultABIInfo {
+public:
+  BPFABIInfo(CodeGenTypes &CGT) : DefaultABIInfo(CGT) {}
+
+  ABIArgInfo classifyReturnType(QualType RetTy) const {
+    if (RetTy->isVoidType())
+      return ABIArgInfo::getIgnore();
+
+    if (isAggregateTypeForABI(RetTy))
+      return getNaturalAlignIndirect(RetTy);
+
+    // Treat an enum type as its underlying type.
+    if (const EnumType *EnumTy = RetTy->getAs<EnumType>())
+      RetTy = EnumTy->getDecl()->getIntegerType();
+
+    ASTContext &Context = getContext();
+    if (const auto *EIT = RetTy->getAs<BitIntType>())
+      if (EIT->getNumBits() > Context.getTypeSize(Context.Int128Ty))
+        return getNaturalAlignIndirect(RetTy);
+
+    // Caller will do necessary sign/zero extension.
+    return ABIArgInfo::getDirect();
+  }
+
+  void computeInfo(CGFunctionInfo &FI) const override {
+    FI.getReturnInfo() = classifyReturnType(FI.getReturnType());
+    for (auto &I : FI.arguments())
+      I.info = classifyArgumentType(I.type);
+  }
+
+};
+
+class BPFTargetCodeGenInfo : public TargetCodeGenInfo {
+public:
+  BPFTargetCodeGenInfo(CodeGenTypes &CGT)
+      : TargetCodeGenInfo(std::make_unique<BPFABIInfo>(CGT)) {}
+
+  const BPFABIInfo &getABIInfo() const {
+    return static_cast<const BPFABIInfo&>(TargetCodeGenInfo::getABIInfo());
+  }
+};
+
+}
+
 //===----------------------------------------------------------------------===//
 // Driver code
 //===----------------------------------------------------------------------===//
@@ -11727,6 +11777,9 @@ const TargetCodeGenInfo &CodeGenModule::getTargetCodeGenInfo() {
                                                       : hasFP64   ? 64
                                                                   : 32));
   }
+  case llvm::Triple::bpfeb:
+  case llvm::Triple::bpfel:
+    return SetCGInfo(new BPFTargetCodeGenInfo(Types));
   }
 }
 

clang/test/CodeGen/bpf-abiinfo.c

@@ -0,0 +1,24 @@
+// REQUIRES: bpf-registered-target
+// RUN: %clang_cc1 -triple bpf -O2 -emit-llvm -disable-llvm-passes %s -o - | FileCheck %s
+
+_Bool bar_bool(void);
+unsigned char bar_char(void);
+short bar_short(void);
+int bar_int(void);
+
+int foo_bool(void) {
+        if (bar_bool() != 1) return 0; else return 1;
+}
+// CHECK: %call = call i1 @bar_bool()
+int foo_char(void) {
+        if (bar_char() != 10) return 0; else return 1;
+}
+// CHECK: %call = call i8 @bar_char()
+int foo_short(void) {
+        if (bar_short() != 10) return 0; else return 1;
+}
+// CHECK: %call = call i16 @bar_short()
+int foo_int(void) {
+        if (bar_int() != 10) return 0; else return 1;
+}
+// CHECK: %call = call i32 @bar_int()