[BPF] Attribute preserve_static_offset for structs

This commit adds a new BPF specific structure attribte
`__attribute__((preserve_static_offset))` and a pass to deal with it.

This attribute may be attached to a struct or union declaration, where
it notifies the compiler that this structure is a "context" structure.
The following limitations apply to context structures:
- runtime environment might patch access to the fields of this type by
  updating the field offset;

  BPF verifier limits access patterns allowed for certain data
  types. E.g. `struct __sk_buff` and `struct bpf_sock_ops`. For these
  types only `LD/ST <reg> <static-offset>` memory loads and stores are
  allowed.

  This is so because offsets of the fields of these structures do not
  match real offsets in the running kernel. During BPF program
  load/verification loads and stores to the fields of these types are
  rewritten so that offsets match real offsets. For this rewrite to
  happen static offsets have to be encoded in the instructions.

  See `kernel/bpf/verifier.c:convert_ctx_access` function in the Linux
  kernel source tree for details.

- runtime environment might disallow access to the field of the type
  through modified pointers.

  During BPF program verification a tag `PTR_TO_CTX` is tracked for
  register values. In case if register with such tag is modified BPF
  programs are not allowed to read or write memory using register. See
  kernel/bpf/verifier.c:check_mem_access function in the Linux kernel
  source tree for details.

Access to the structure fields is translated to IR as a sequence:
- `(load (getelementptr %ptr %offset))` or
- `(store (getelementptr %ptr %offset))`

During instruction selection phase such sequences are translated as a
single load instruction with embedded offset, e.g. `LDW %ptr, %offset`,
which matches access pattern necessary for the restricted
set of types described above (when `%offset` is static).

Multiple optimizer passes might separate these instructions, this
includes:
- SimplifyCFGPass (sinking)
- InstCombine (sinking)
- GVN (hoisting)

The `preserve_static_offset` attribute marks structures for which the
following transformations happen:
- at the early IR processing stage:
  - `(load (getelementptr ...))` replaced by call to intrinsic
    `llvm.bpf.getelementptr.and.load`;
  - `(store (getelementptr ...))` replaced by call to intrinsic
    `llvm.bpf.getelementptr.and.store`;
- at the late IR processing stage this modification is undone.

Such handling prevents various optimizer passes from generating
sequences of instructions that would be rejected by BPF verifier.

The __attribute__((preserve_static_offset)) has a priority over
__attribute__((preserve_access_index)). When preserve_access_index
attribute is present preserve access index transformations are not
applied.

This addresses the issue reported by the following thread:

https://lore.kernel.org/bpf/CAA-VZPmxh8o8EBcJ=m-DH4ytcxDFmo0JKsm1p1gf40kS0CE3NQ@mail.gmail.com/T/#m4b9ce2ce73b34f34172328f975235fc6f19841b6

This is a second attempt to commit this change, previous reverted
commit is: cb13e92.
The following items had been fixed:
- test case bpf-preserve-static-offset-bitfield.c now uses
  `-triple bpfel` to avoid different codegen for little/big endian
  targets.
- BPFPreserveStaticOffset.cpp:removePAICalls() modified to avoid
  use after free for `WorkList` elements `V`.

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

Author: eddyz87

clang/include/clang/Basic/Attr.td

@@ -2024,6 +2024,14 @@ def BPFPreserveAccessIndex : InheritableAttr,
   let LangOpts = [COnly];
 }
 
+def BPFPreserveStaticOffset : InheritableAttr,
+                              TargetSpecificAttr<TargetBPF>  {
+  let Spellings = [Clang<"preserve_static_offset">];
+  let Subjects = SubjectList<[Record], ErrorDiag>;
+  let Documentation = [BPFPreserveStaticOffsetDocs];
+  let LangOpts = [COnly];
+}
+
 def BTFDeclTag : InheritableAttr {
   let Spellings = [Clang<"btf_decl_tag">];
   let Args = [StringArgument<"BTFDeclTag">];

clang/include/clang/Basic/AttrDocs.td

@@ -2199,6 +2199,43 @@ preserving struct or union member access debuginfo indices of this
 struct or union, similar to clang ``__builtin_preserve_access_index()``.
   }];
 }
+
+def BPFPreserveStaticOffsetDocs : Documentation {
+  let Category = DocCatFunction;
+  let Content = [{
+Clang supports the ``__attribute__((preserve_static_offset))``
+attribute for the BPF target. This attribute may be attached to a
+struct or union declaration. Reading or writing fields of types having
+such annotation is guaranteed to generate LDX/ST/STX instruction with
+offset corresponding to the field.
+
+For example:
+
+.. code-block:: c
+
+  struct foo {
+    int a;
+    int b;
+  };
+
+  struct bar {
+    int a;
+    struct foo b;
+  } __attribute__((preserve_static_offset));
+
+  void buz(struct bar *g) {
+    g->b.a = 42;
+  }
+
+The assignment to ``g``'s field would produce an ST instruction with
+offset 8: ``*(u32)(r1 + 8) = 42;``.
+
+Without this attribute generated instructions might be different,
+depending on optimizations behavior. E.g. the example above could be
+rewritten as ``r1 += 8; *(u32)(r1 + 0) = 42;``.
+  }];
+}
+
 def BTFDeclTagDocs : Documentation {
   let Category = DocCatFunction;
   let Content = [{

clang/lib/CodeGen/CGExpr.cpp

@@ -3833,6 +3833,33 @@ static QualType getFixedSizeElementType(const ASTContext &ctx,
   return eltType;
 }
 
+static bool hasBPFPreserveStaticOffset(const RecordDecl *D) {
+  return D && D->hasAttr<BPFPreserveStaticOffsetAttr>();
+}
+
+static bool hasBPFPreserveStaticOffset(const Expr *E) {
+  if (!E)
+    return false;
+  QualType PointeeType = E->getType()->getPointeeType();
+  if (PointeeType.isNull())
+    return false;
+  if (const auto *BaseDecl = PointeeType->getAsRecordDecl())
+    return hasBPFPreserveStaticOffset(BaseDecl);
+  return false;
+}
+
+// Wraps Addr with a call to llvm.preserve.static.offset intrinsic.
+static Address wrapWithBPFPreserveStaticOffset(CodeGenFunction &CGF,
+                                               Address &Addr) {
+  if (!CGF.getTarget().getTriple().isBPF())
+    return Addr;
+
+  llvm::Function *Fn =
+      CGF.CGM.getIntrinsic(llvm::Intrinsic::preserve_static_offset);
+  llvm::CallInst *Call = CGF.Builder.CreateCall(Fn, {Addr.getPointer()});
+  return Address(Call, Addr.getElementType(), Addr.getAlignment());
+}
+
 /// Given an array base, check whether its member access belongs to a record
 /// with preserve_access_index attribute or not.
 static bool IsPreserveAIArrayBase(CodeGenFunction &CGF, const Expr *ArrayBase) {
@@ -3894,6 +3921,9 @@ static Address emitArraySubscriptGEP(CodeGenFunction &CGF, Address addr,
   CharUnits eltAlign =
     getArrayElementAlign(addr.getAlignment(), indices.back(), eltSize);
 
+  if (hasBPFPreserveStaticOffset(Base))
+    addr = wrapWithBPFPreserveStaticOffset(CGF, addr);
+
   llvm::Value *eltPtr;
   auto LastIndex = dyn_cast<llvm::ConstantInt>(indices.back());
   if (!LastIndex ||
@@ -4522,6 +4552,8 @@ LValue CodeGenFunction::EmitLValueForField(LValue base,
     Address Addr = base.getAddress(*this);
     unsigned Idx = RL.getLLVMFieldNo(field);
     const RecordDecl *rec = field->getParent();
+    if (hasBPFPreserveStaticOffset(rec))
+      Addr = wrapWithBPFPreserveStaticOffset(*this, Addr);
     if (!UseVolatile) {
       if (!IsInPreservedAIRegion &&
           (!getDebugInfo() || !rec->hasAttr<BPFPreserveAccessIndexAttr>())) {
@@ -4594,6 +4626,8 @@ LValue CodeGenFunction::EmitLValueForField(LValue base,
   }
 
   Address addr = base.getAddress(*this);
+  if (hasBPFPreserveStaticOffset(rec))
+    addr = wrapWithBPFPreserveStaticOffset(*this, addr);
   if (auto *ClassDef = dyn_cast<CXXRecordDecl>(rec)) {
     if (CGM.getCodeGenOpts().StrictVTablePointers &&
         ClassDef->isDynamicClass()) {

clang/lib/Sema/SemaDeclAttr.cpp

@@ -9036,6 +9036,9 @@ ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D, const ParsedAttr &AL,
   case ParsedAttr::AT_BPFPreserveAccessIndex:
     handleBPFPreserveAccessIndexAttr(S, D, AL);
     break;
+  case ParsedAttr::AT_BPFPreserveStaticOffset:
+    handleSimpleAttribute<BPFPreserveStaticOffsetAttr>(S, D, AL);
+    break;
   case ParsedAttr::AT_BTFDeclTag:
     handleBTFDeclTagAttr(S, D, AL);
     break;

clang/test/CodeGen/bpf-preserve-static-offset-arr.c

@@ -0,0 +1,33 @@
+// NOTE: Assertions have been autogenerated by utils/update_cc_test_checks.py UTC_ARGS: --version 2
+// REQUIRES: bpf-registered-target
+// RUN: %clang -cc1 -triple bpf -disable-llvm-passes -S -emit-llvm -o - %s \
+// RUN: | FileCheck %s
+
+// Check that call to preserve.static.offset is generated when array
+// member of a struct marked with __attribute__((preserve_static_offset))
+// is accessed.
+
+#define __ctx __attribute__((preserve_static_offset))
+
+struct foo {
+  struct {
+    int a;
+  } b[7];
+} __ctx;
+
+// CHECK-LABEL: define dso_local i32 @arr_access
+// CHECK-SAME: (ptr noundef [[P:%.*]]) #[[ATTR0:[0-9]+]] {
+// CHECK-NEXT:  entry:
+// CHECK-NEXT:    [[P_ADDR:%.*]] = alloca ptr, align 8
+// CHECK-NEXT:    store ptr [[P]], ptr [[P_ADDR]], align 8
+// CHECK-NEXT:    [[TMP0:%.*]] = load ptr, ptr [[P_ADDR]], align 8
+// CHECK-NEXT:    [[TMP1:%.*]] = call ptr @llvm.preserve.static.offset(ptr [[TMP0]])
+// CHECK-NEXT:    [[B:%.*]] = getelementptr inbounds [[STRUCT_FOO:%.*]], ptr [[TMP1]], i32 0, i32 0
+// CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds [7 x %struct.anon], ptr [[B]], i64 0, i64 2
+// CHECK-NEXT:    [[A:%.*]] = getelementptr inbounds [[STRUCT_ANON:%.*]], ptr [[ARRAYIDX]], i32 0, i32 0
+// CHECK-NEXT:    [[TMP2:%.*]] = load i32, ptr [[A]], align 4
+// CHECK-NEXT:    ret i32 [[TMP2]]
+//
+int arr_access(struct foo *p) {
+  return p->b[2].a;
+}

clang/test/CodeGen/bpf-preserve-static-offset-bitfield.c

@@ -0,0 +1,31 @@
+// NOTE: Assertions have been autogenerated by utils/update_cc_test_checks.py UTC_ARGS: --version 2
+// REQUIRES: bpf-registered-target
+// RUN: %clang -cc1 -triple bpfel -disable-llvm-passes -S -emit-llvm -o - %s \
+// RUN: | FileCheck %s
+
+// Check that call to preserve.static.offset is generated when bitfield
+// from a struct marked with __attribute__((preserve_static_offset)) is
+// accessed.
+
+#define __ctx __attribute__((preserve_static_offset))
+
+struct foo {
+  unsigned a:1;
+} __ctx;
+
+// CHECK-LABEL: define dso_local void @lvalue_bitfield
+// CHECK-SAME: (ptr noundef [[P:%.*]]) #[[ATTR0:[0-9]+]] {
+// CHECK-NEXT:  entry:
+// CHECK-NEXT:    [[P_ADDR:%.*]] = alloca ptr, align 8
+// CHECK-NEXT:    store ptr [[P]], ptr [[P_ADDR]], align 8
+// CHECK-NEXT:    [[TMP0:%.*]] = load ptr, ptr [[P_ADDR]], align 8
+// CHECK-NEXT:    [[TMP1:%.*]] = call ptr @llvm.preserve.static.offset(ptr [[TMP0]])
+// CHECK-NEXT:    [[BF_LOAD:%.*]] = load i8, ptr [[TMP1]], align 4
+// CHECK-NEXT:    [[BF_CLEAR:%.*]] = and i8 [[BF_LOAD]], -2
+// CHECK-NEXT:    [[BF_SET:%.*]] = or i8 [[BF_CLEAR]], 1
+// CHECK-NEXT:    store i8 [[BF_SET]], ptr [[TMP1]], align 4
+// CHECK-NEXT:    ret void
+//
+void lvalue_bitfield(struct foo *p) {
+  p->a = 1;
+}

clang/test/CodeGen/bpf-preserve-static-offset-lvalue.c

@@ -0,0 +1,28 @@
+// NOTE: Assertions have been autogenerated by utils/update_cc_test_checks.py UTC_ARGS: --version 2
+// REQUIRES: bpf-registered-target
+// RUN: %clang -cc1 -triple bpf -disable-llvm-passes -S -emit-llvm -o - %s \
+// RUN: | FileCheck %s
+
+// Check that call to preserve.static.offset is generated when field of
+// a struct marked with __attribute__((preserve_static_offset)) is accessed.
+
+#define __ctx __attribute__((preserve_static_offset))
+
+struct foo {
+  int a;
+} __ctx;
+
+// CHECK-LABEL: define dso_local void @lvalue
+// CHECK-SAME: (ptr noundef [[P:%.*]]) #[[ATTR0:[0-9]+]] {
+// CHECK-NEXT:  entry:
+// CHECK-NEXT:    [[P_ADDR:%.*]] = alloca ptr, align 8
+// CHECK-NEXT:    store ptr [[P]], ptr [[P_ADDR]], align 8
+// CHECK-NEXT:    [[TMP0:%.*]] = load ptr, ptr [[P_ADDR]], align 8
+// CHECK-NEXT:    [[TMP1:%.*]] = call ptr @llvm.preserve.static.offset(ptr [[TMP0]])
+// CHECK-NEXT:    [[A:%.*]] = getelementptr inbounds [[STRUCT_FOO:%.*]], ptr [[TMP1]], i32 0, i32 0
+// CHECK-NEXT:    store i32 42, ptr [[A]], align 4
+// CHECK-NEXT:    ret void
+//
+void lvalue(struct foo *p) {
+  p->a = 42;
+}

clang/test/CodeGen/bpf-preserve-static-offset-non-bpf.c

@@ -0,0 +1,18 @@
+// REQUIRES: x86-registered-target
+// RUN: %clang -cc1 -triple x86_64 -disable-llvm-passes -S -emit-llvm  -o - %s \
+// RUN: | FileCheck %s
+
+// Verify that __attribute__((preserve_static_offset))
+// has no effect for non-BPF target.
+
+#define __ctx __attribute__((preserve_static_offset))
+
+struct foo {
+  int a;
+} __ctx;
+
+// CHECK-NOT: @llvm_preserve_static_offset
+
+int bar(struct foo *p) {
+  return p->a;
+}

clang/test/CodeGen/bpf-preserve-static-offset-pai.c

@@ -0,0 +1,29 @@
+// NOTE: Assertions have been autogenerated by utils/update_cc_test_checks.py UTC_ARGS: --version 2
+// REQUIRES: bpf-registered-target
+// RUN: %clang -cc1 -triple bpf -disable-llvm-passes -S -emit-llvm -o - %s \
+// RUN: | FileCheck %s
+
+// Verify that preserve_static_offset does not interfere with
+// preserve_access_index at IR generation stage.
+
+#define __ctx __attribute__((preserve_static_offset))
+#define __pai __attribute__((preserve_access_index))
+
+struct foo {
+  int a;
+} __ctx __pai;
+
+// CHECK-LABEL: define dso_local i32 @bar
+// CHECK-SAME: (ptr noundef [[P:%.*]]) #[[ATTR0:[0-9]+]] {
+// CHECK-NEXT:  entry:
+// CHECK-NEXT:    [[P_ADDR:%.*]] = alloca ptr, align 8
+// CHECK-NEXT:    store ptr [[P]], ptr [[P_ADDR]], align 8
+// CHECK-NEXT:    [[TMP0:%.*]] = load ptr, ptr [[P_ADDR]], align 8
+// CHECK-NEXT:    [[TMP1:%.*]] = call ptr @llvm.preserve.static.offset(ptr [[TMP0]])
+// CHECK-NEXT:    [[A:%.*]] = getelementptr inbounds [[STRUCT_FOO:%.*]], ptr [[TMP1]], i32 0, i32 0
+// CHECK-NEXT:    [[TMP2:%.*]] = load i32, ptr [[A]], align 4
+// CHECK-NEXT:    ret i32 [[TMP2]]
+//
+int bar(struct foo *p) {
+  return p->a;
+}

clang/test/Misc/pragma-attribute-supported-attributes-list.test

@@ -23,6 +23,7 @@
 // CHECK-NEXT: Availability ((SubjectMatchRule_record, SubjectMatchRule_enum, SubjectMatchRule_enum_constant, SubjectMatchRule_field, SubjectMatchRule_function, SubjectMatchRule_namespace, SubjectMatchRule_objc_category, SubjectMatchRule_objc_implementation, SubjectMatchRule_objc_interface, SubjectMatchRule_objc_method, SubjectMatchRule_objc_property, SubjectMatchRule_objc_protocol, SubjectMatchRule_record, SubjectMatchRule_type_alias, SubjectMatchRule_variable))
 // CHECK-NEXT: AvailableOnlyInDefaultEvalMethod (SubjectMatchRule_type_alias)
 // CHECK-NEXT: BPFPreserveAccessIndex (SubjectMatchRule_record)
+// CHECK-NEXT: BPFPreserveStaticOffset (SubjectMatchRule_record)
 // CHECK-NEXT: BTFDeclTag (SubjectMatchRule_variable, SubjectMatchRule_function, SubjectMatchRule_record, SubjectMatchRule_field, SubjectMatchRule_type_alias)
 // CHECK-NEXT: BuiltinAlias (SubjectMatchRule_function)
 // CHECK-NEXT: CFAuditedTransfer (SubjectMatchRule_function)

clang/test/Sema/bpf-attr-preserve-static-offset-warns-nonbpf.c

@@ -0,0 +1,6 @@
+// RUN: %clang_cc1 -fsyntax-only -verify %s
+
+#define __pso __attribute__((preserve_static_offset))
+
+struct foo { int a; } __pso; // expected-warning{{unknown attribute}}
+union quux { int a; } __pso; // expected-warning{{unknown attribute}}

clang/test/Sema/bpf-attr-preserve-static-offset-warns.c

@@ -0,0 +1,23 @@
+// RUN: %clang_cc1 -fsyntax-only -verify -triple bpf-pc-linux-gnu %s
+
+#define __pso __attribute__((preserve_static_offset))
+
+// These are correct usages.
+struct foo { int a; } __pso;
+union quux { int a; } __pso;
+struct doug { int a; } __pso __attribute__((packed));
+
+// Rest are incorrect usages.
+typedef int bar __pso;    // expected-error{{attribute only applies to}}
+struct goo {
+  int a __pso;            // expected-error{{attribute only applies to}}
+};
+int g __pso;              // expected-error{{attribute only applies to}}
+__pso void ffunc1(void);  // expected-error{{attribute only applies to}}
+void ffunc2(int a __pso); // expected-error{{attribute only applies to}}
+void ffunc3(void) {
+  int a __pso;            // expected-error{{attribute only applies to}}
+}
+
+struct buz { int a; } __attribute__((preserve_static_offset("hello"))); // \
+  expected-error{{attribute takes no arguments}}

clang/test/Sema/bpf-attr-preserve-static-offset.c

@@ -0,0 +1,27 @@
+// RUN: %clang_cc1 -fsyntax-only -ast-dump -triple bpf-pc-linux-gnu %s | FileCheck %s
+
+// The 'preserve_static_offset' attribute should be propagated to
+// inline declarations (foo's 'b', 'bb', 'c' but not 'd').
+//
+// CHECK:      RecordDecl {{.*}} struct foo definition
+// CHECK-NEXT:   BPFPreserveStaticOffsetAttr
+// CHECK-NEXT:   FieldDecl {{.*}} a
+// CHECK-NEXT:   RecordDecl {{.*}} struct definition
+// CHECK-NEXT:     FieldDecl {{.*}} aa
+// CHECK-NEXT:   FieldDecl {{.*}} b
+// CHECK-NEXT: RecordDecl {{.*}} union bar definition
+// CHECK-NEXT:   BPFPreserveStaticOffsetAttr
+// CHECK-NEXT:   FieldDecl {{.*}} a
+// CHECK-NEXT:   FieldDecl {{.*}} b
+
+struct foo {
+  int a;
+  struct {
+    int aa;
+  } b;
+} __attribute__((preserve_static_offset));
+
+union bar {
+  int a;
+  long b;
+} __attribute__((preserve_static_offset));

llvm/include/llvm/IR/Intrinsics.td

@@ -2469,6 +2469,10 @@ def int_preserve_struct_access_index : DefaultAttrsIntrinsic<[llvm_anyptr_ty],
                                                  [IntrNoMem,
                                                   ImmArg<ArgIndex<1>>,
                                                   ImmArg<ArgIndex<2>>]>;
+def int_preserve_static_offset : DefaultAttrsIntrinsic<[llvm_ptr_ty],
+                                                       [llvm_ptr_ty],
+                                                       [IntrNoMem, IntrSpeculatable,
+                                                        ReadNone <ArgIndex<0>>]>;
 
 //===------------ Intrinsics to perform common vector shuffles ------------===//
 

llvm/include/llvm/IR/IntrinsicsBPF.td

@@ -37,4 +37,43 @@ let TargetPrefix = "bpf" in {  // All intrinsics start with "llvm.bpf."
   def int_bpf_compare : ClangBuiltin<"__builtin_bpf_compare">,
               Intrinsic<[llvm_i1_ty], [llvm_i32_ty, llvm_anyint_ty, llvm_anyint_ty],
               [IntrNoMem]>;
+  def int_bpf_getelementptr_and_load : ClangBuiltin<"__builtin_bpf_getelementptr_and_load">,
+              Intrinsic<[llvm_any_ty],
+                        [llvm_ptr_ty,     // base ptr for getelementptr
+                         llvm_i1_ty,      // volatile
+                         llvm_i8_ty,      // atomic order
+                         llvm_i8_ty,      // synscope id
+                         llvm_i8_ty,      // alignment
+                         llvm_i1_ty,      // inbounds
+                         llvm_vararg_ty], // indices for getelementptr insn
+                        [IntrNoCallback,
+                         IntrNoFree,
+                         IntrWillReturn,
+                         NoCapture <ArgIndex<0>>,
+                         ImmArg    <ArgIndex<1>>, // volatile
+                         ImmArg    <ArgIndex<2>>, // atomic order
+                         ImmArg    <ArgIndex<3>>, // synscope id
+                         ImmArg    <ArgIndex<4>>, // alignment
+                         ImmArg    <ArgIndex<5>>, // inbounds
+                        ]>;
+  def int_bpf_getelementptr_and_store : ClangBuiltin<"__builtin_bpf_getelementptr_and_store">,
+              Intrinsic<[],
+                        [llvm_any_ty,     // value to store
+                         llvm_ptr_ty,     // base ptr for getelementptr
+                         llvm_i1_ty,      // volatile
+                         llvm_i8_ty,      // atomic order
+                         llvm_i8_ty,      // syncscope id
+                         llvm_i8_ty,      // alignment
+                         llvm_i1_ty,      // inbounds
+                         llvm_vararg_ty], // indexes for getelementptr insn
+                        [IntrNoCallback,
+                         IntrNoFree,
+                         IntrWillReturn,
+                         NoCapture <ArgIndex<1>>,
+                         ImmArg    <ArgIndex<2>>, // volatile
+                         ImmArg    <ArgIndex<3>>, // atomic order
+                         ImmArg    <ArgIndex<4>>, // syncscope id
+                         ImmArg    <ArgIndex<5>>, // alignment
+                         ImmArg    <ArgIndex<6>>, // inbounds
+                        ]>;
 }