[SanitizerBinaryMetadata] Emit constants as ULEB128

[melver]

Summary

Emit all constant integers produced by SanitizerBinaryMetadata as
ULEB128 to further reduce binary space used. Increasing the version is
not necessary given this change depends on (and will land) along with
the bump to v2.

To support this, the !pcsections metadata format is extended to allow
for per-section options, encoded in the first MD operator which must
always be a string and contain the section: "<section>!<options>".

Depends on D143482

[melver]

Move AsmPrinter LEB128 helpers to AsmPrinter.cpp and use them.

clang/test/CodeGen/sanitize-metadata.c

 // ATOMICS-LABEL: __sanitizer_metadata_atomics.module_dtor
 // ATOMICS: call void @__sanitizer_metadata_atomics_del(i32 2, ptr @__start_sanmd_atomics, ptr @__stop_sanmd_atomics)
 
 // CHECK-LABEL: __sanitizer_metadata_covered.module_ctor
 // CHECK: call void @__sanitizer_metadata_covered_add(i32 2, ptr @__start_sanmd_covered, ptr @__stop_sanmd_covered)
 // CHECK-LABEL: __sanitizer_metadata_covered.module_dtor
 // CHECK: call void @__sanitizer_metadata_covered_del(i32 2, ptr @__start_sanmd_covered, ptr @__stop_sanmd_covered)
 
-// ATOMICS: ![[ATOMICS_COVERED]] = !{!"sanmd_covered", ![[ATOMICS_COVERED_AUX:[0-9]+]]}
-// ATOMICS: ![[ATOMICS_COVERED_AUX]] = !{i8 1}
-// ATOMICS: ![[ATOMIC_OP]] = !{!"sanmd_atomics"}
+// ATOMICS: ![[ATOMICS_COVERED]] = !{!"sanmd_covered!C", ![[ATOMICS_COVERED_AUX:[0-9]+]]}
+// ATOMICS: ![[ATOMICS_COVERED_AUX]] = !{i64 1}
+// ATOMICS: ![[ATOMIC_OP]] = !{!"sanmd_atomics!C"}

compiler-rt/test/metadata/common.h

template <typename T> T consume(const char *&pos, const char *end) {
   // We need to memcpy from pos, because it's not guaranteed that every entry
   // has the required alignment of T.
   memcpy(&v, pos, sizeof(T));
   pos += sizeof(T);
   assert(pos <= end);
   return v;
 }
 
+uint64_t consume_uleb128(const char *&pos, const char *end) {
+  uint64_t val = 0;
+  int shift = 0;
+  uint8_t cur;
+  do {
+    cur = *pos++;
+    val |= uint64_t{cur & 0x7fu} << shift;
+    shift += 7;
+  } while (cur & 0x80);
+  assert(shift < 64);
+  assert(pos <= end);
+  return val;
+}
+
 constexpr uint32_t kSanitizerBinaryMetadataUARHasSize = 1 << 2;
 
 uint32_t meta_version;
 const char *meta_start;
 const char *meta_end;
 const char *atomics_start;
 const char *atomics_end;
 }; // namespace
 
 extern "C" {
 void __sanitizer_metadata_covered_add(uint32_t version, const char *start,
                                       const char *end) {
   const uint32_t version_base = version & 0xffff;
   const bool offset_ptr_sized = version & (1 << 16);
   assert(version_base == 2);
   printf("metadata add version %u\n", version_base);
   for (const char *pos = start; pos < end;) {
     const auto base = reinterpret_cast<uintptr_t>(pos);
     const intptr_t offset = offset_ptr_sized ? consume<intptr_t>(pos, end)
                                              : consume<int32_t>(pos, end);
-    [[maybe_unused]] const uint32_t size = consume<uint32_t>(pos, end);
-    const uint32_t features = consume<uint8_t>(pos, end);
-    uint32_t stack_args = 0;
+    [[maybe_unused]] const uint64_t size = consume_uleb128(pos, end);
+    const uint64_t features = consume_uleb128(pos, end);
+    uint64_t stack_args = 0;
     if (features & kSanitizerBinaryMetadataUARHasSize)
-      stack_args = consume<uint32_t>(pos, end);
+      stack_args = consume_uleb128(pos, end);
     if (const char *name = symbolize(base + offset))
-      printf("%s: features=%x stack_args=%u\n", name, features, stack_args);
+      printf("%s: features=%lx stack_args=%lu\n", name, features, stack_args);
   }
   meta_version = version;
   meta_start = start;
   meta_end = end;
 }
 
 void __sanitizer_metadata_covered_del(uint32_t version, const char *start,
                                       const char *end) {

llvm/docs/PCSectionsMetadata.rst

 To avoid relocations in the final binary, each PC address stored at ``entry``
 is a relative relocation, computed as ``pc - entry``. To decode, a user has to
 compute ``entry + *entry``.
 
 The size of each entry depends on the code model. With large and medium sized
 code models, the entry size matches pointer size. For any smaller code model
 the entry size is just 32 bits.
 
+Encoding Options
+----------------
+
+Optional encoding options can be passed in the first ``MDString`` operator:
+``<section>!<options>``. The following options are available:
+
+    * ``C`` -- Compress constant integers of size 2-8 bytes as ULEB128; this
+      includes the function size (but excludes the PC entry).
+
+For example, ``foo!C`` will emit into section ``foo`` with all constants
+encoded as ULEB128.
+
 Guarantees on Code Generation
 =============================
 
 Attaching ``!pcsections`` metadata to LLVM IR instructions *shall not* affect
 optimizations or code generation outside the requested PC sections.
 
 While relying on LLVM IR metadata to request PC sections makes the above
 guarantee relatively trivial, propagation of metadata through the optimization

llvm/include/llvm/Transforms/Instrumentation/SanitizerBinaryMetadata.h

struct SanitizerBinaryMetadataOptions {
   bool UAR = false;
   SanitizerBinaryMetadataOptions() = default;
 };
 
 inline constexpr int kSanitizerBinaryMetadataAtomicsBit = 0;
 inline constexpr int kSanitizerBinaryMetadataUARBit = 1;
 inline constexpr int kSanitizerBinaryMetadataUARHasSizeBit = 2;
 
-inline constexpr uint32_t kSanitizerBinaryMetadataAtomics =
+inline constexpr uint64_t kSanitizerBinaryMetadataAtomics =
     1 << kSanitizerBinaryMetadataAtomicsBit;
-inline constexpr uint32_t kSanitizerBinaryMetadataUAR =
+inline constexpr uint64_t kSanitizerBinaryMetadataUAR =
     1 << kSanitizerBinaryMetadataUARBit;
-inline constexpr uint32_t kSanitizerBinaryMetadataUARHasSize =
+inline constexpr uint64_t kSanitizerBinaryMetadataUARHasSize =
     1 << kSanitizerBinaryMetadataUARHasSizeBit;
 
 inline constexpr char kSanitizerBinaryMetadataCoveredSection[] =
     "sanmd_covered";
 inline constexpr char kSanitizerBinaryMetadataAtomicsSection[] =
     "sanmd_atomics";
 
 /// Public interface to the SanitizerBinaryMetadata module pass for emitting

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

void AsmPrinter::emitPCSections(const MachineFunction &MF) {
   };
   // Emit symbols into sections and data as specified in the pcsections MDNode.
   auto EmitForMD = [&](const MDNode &MD, ArrayRef<const MCSymbol *> Syms,
                        bool Deltas) {
     // Expect the first operand to be a section name. After that, a tuple of
     // constants may appear, which will simply be emitted into the current
     // section (the user of MD_pcsections decides the format of encoded data).
     assert(isa<MDString>(MD.getOperand(0)) && "first operand not a string");
+    bool ConstULEB128 = false;
     for (const MDOperand &MDO : MD.operands()) {
       if (auto *S = dyn_cast<MDString>(MDO)) {
-        SwitchSection(S->getString());
+        // Found string, start of new section!
+        // Find options for this section "<section>!<opts>" - supported options:
+        //   C = Compress constant integers of size 2-8 bytes as ULEB128.
+        const StringRef SecWithOpt = S->getString();
+        const size_t OptStart = SecWithOpt.find('!'); // likely npos
+        const StringRef Sec = SecWithOpt.substr(0, OptStart);
+        const StringRef Opts = SecWithOpt.substr(OptStart); // likely empty
+        ConstULEB128 = Opts.find('C') != StringRef::npos;
+#ifndef NDEBUG
+        for (char O : Opts)
+          assert((O == '!' || O == 'C') && "Invalid !pcsections options");
+#endif
+        SwitchSection(Sec);
         const MCSymbol *Prev = Syms.front();
         for (const MCSymbol *Sym : Syms) {
           if (Sym == Prev || !Deltas) {
             // Use the entry itself as the base of the relative offset.
             MCSymbol *Base = MF.getContext().createTempSymbol("pcsection_base");
             OutStreamer->emitLabel(Base);
             // Emit relative relocation `addr - base`, which avoids a dynamic
             // relocation in the final binary. User will get the address with
             // `base + addr`.
             emitLabelDifference(Sym, Base, RelativeRelocSize);
           } else {
+            // Emit delta between symbol and previous symbol.
+            if (ConstULEB128)
+              OutStreamer->emitAbsoluteSymbolDiffAsULEB128(Sym, Prev);
+            else
               emitLabelDifference(Sym, Prev, 4);
           }
           Prev = Sym;
         }
       } else {
+        // Emit auxiliary data after PC.
         assert(isa<MDNode>(MDO) && "expecting either string or tuple");
         const auto *AuxMDs = cast<MDNode>(MDO);
         for (const MDOperand &AuxMDO : AuxMDs->operands()) {
           assert(isa<ConstantAsMetadata>(AuxMDO) && "expecting a constant");
-          const auto *C = cast<ConstantAsMetadata>(AuxMDO);
-          emitGlobalConstant(F.getParent()->getDataLayout(), C->getValue());
+          const Constant *C = cast<ConstantAsMetadata>(AuxMDO)->getValue();
+          const DataLayout &DL = F.getParent()->getDataLayout();
+          const uint64_t Size = DL.getTypeStoreSize(C->getType());
+
+          if (auto *CI = dyn_cast<ConstantInt>(C);
+              CI && ConstULEB128 && Size > 1 && Size <= 8) {
+            OutStreamer->emitULEB128IntValue(CI->getZExtValue());
+          } else {
+            emitGlobalConstant(DL, C);
+          }
         }
       }
     }
   };
 
   OutStreamer->pushSection();
   // Emit PCs for function start and function size.
   if (const MDNode *MD = F.getMetadata(LLVMContext::MD_pcsections))

llvm/lib/CodeGen/SanitizerBinaryMetadata.cpp

initializeMachineSanitizerBinaryMetadataPass(
       *PassRegistry::getPassRegistry());
 }
 
 bool MachineSanitizerBinaryMetadata::runOnMachineFunction(MachineFunction &MF) {
   MDNode *MD = MF.getFunction().getMetadata(LLVMContext::MD_pcsections);
   if (!MD)
     return false;
   const auto &Section = *cast<MDString>(MD->getOperand(0));
-  if (!Section.getString().equals(kSanitizerBinaryMetadataCoveredSection))
+  if (!Section.getString().startswith(kSanitizerBinaryMetadataCoveredSection))
     return false;
   auto &AuxMDs = *cast<MDTuple>(MD->getOperand(1));
   // Assume it currently only has features.
   assert(AuxMDs.getNumOperands() == 1);
   Constant *Features =
       cast<ConstantAsMetadata>(AuxMDs.getOperand(0))->getValue();
   if (!Features->getUniqueInteger()[kSanitizerBinaryMetadataUARBit])
     return false;

llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp

 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Value.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/ProfileData/InstrProf.h"
+#include "llvm/Support/Allocator.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/StringSaver.h"
 #include "llvm/Transforms/Instrumentation.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
 #include <array>
 #include <cstdint>
-#include <limits>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "sanmd"
 
 namespace {
 
 //===--- Constants --------------------------------------------------------===//
private:
   // Determines which set of metadata to collect for this instruction.
   //
   // Returns true if covered metadata is required to unambiguously interpret
   // other metadata. For example, if we are interested in atomics metadata, any
   // function with memory operations (atomic or not) requires covered metadata
   // to determine if a memory operation is atomic or not in modules compiled
   // with SanitizerBinaryMetadata.
   bool runOn(Instruction &I, MetadataInfoSet &MIS, MDBuilder &MDB,
-             uint32_t &FeatureMask);
+             uint64_t &FeatureMask);
 
   // Get start/end section marker pointer.
   GlobalVariable *getSectionMarker(const Twine &MarkerName, Type *Ty);
 
   // Returns the target-dependent section name.
   StringRef getSectionName(StringRef SectionSuffix);
 
   // Returns the section start marker name.
   Twine getSectionStart(StringRef SectionSuffix);
 
   // Returns the section end marker name.
   Twine getSectionEnd(StringRef SectionSuffix);
 
   // Returns true if the access to the address should be considered "atomic".
   bool pretendAtomicAccess(const Value *Addr);
 
   Module &Mod;
   const SanitizerBinaryMetadataOptions Options;
   const Triple TargetTriple;
   IRBuilder<> IRB;
+  BumpPtrAllocator Alloc;
+  UniqueStringSaver StringPool{Alloc};
 };
 
 bool SanitizerBinaryMetadata::run() {
   MetadataInfoSet MIS;
 
   for (Function &F : Mod)
     runOn(F, MIS);
 
void SanitizerBinaryMetadata::runOn(Function &F, MetadataInfoSet &MIS) {
   // Don't touch available_externally functions, their actual body is elsewhere.
   if (F.getLinkage() == GlobalValue::AvailableExternallyLinkage)
     return;
 
   MDBuilder MDB(F.getContext());
 
   // The metadata features enabled for this function, stored along covered
   // metadata (if enabled).
-  uint32_t FeatureMask = 0;
+  uint64_t FeatureMask = 0;
   // Don't emit unnecessary covered metadata for all functions to save space.
   bool RequiresCovered = false;
 
   if (Options.Atomics || Options.UAR) {
     for (BasicBlock &BB : F)
       for (Instruction &I : BB)
         RequiresCovered |= runOn(I, MIS, MDB, FeatureMask);
   }
 
   if (F.isVarArg())
     FeatureMask &= ~kSanitizerBinaryMetadataUAR;
   if (FeatureMask & kSanitizerBinaryMetadataUAR) {
     RequiresCovered = true;
     NumMetadataUAR++;
   }
 
   // Covered metadata is always emitted if explicitly requested, otherwise only
   // if some other metadata requires it to unambiguously interpret it for
   // modules compiled with SanitizerBinaryMetadata.
   if (Options.Covered || (FeatureMask && RequiresCovered)) {
     NumMetadataCovered++;
     const auto *MI = &MetadataInfo::Covered;
     MIS.insert(MI);
     const StringRef Section = getSectionName(MI->SectionSuffix);
-    // The feature mask will be placed after the size of the function.
-    assert(FeatureMask <= std::numeric_limits<uint8_t>::max() &&
-           "Increase feature mask bytes and bump version");
-    Constant *CFM = IRB.getInt8(FeatureMask);
+    // The feature mask will be placed after the function size.
+    Constant *CFM = IRB.getInt64(FeatureMask);
     F.setMetadata(LLVMContext::MD_pcsections,
                   MDB.createPCSections({{Section, {CFM}}}));
   }
 }
 
 bool isUARSafeCall(CallInst *CI) {
   auto *F = CI->getCalledFunction();
   // There are no intrinsic functions that leak arguments.
if (auto *GV = dyn_cast<GlobalVariable>(Addr)) {
     if (GV->isConstant())
       return false; // Shared, but not mutable.
   }
 
   return true;
 }
 
 bool SanitizerBinaryMetadata::runOn(Instruction &I, MetadataInfoSet &MIS,
-                                    MDBuilder &MDB, uint32_t &FeatureMask) {
+                                    MDBuilder &MDB, uint64_t &FeatureMask) {
   SmallVector<const MetadataInfo *, 1> InstMetadata;
   bool RequiresCovered = false;
 
   // Only call if at least 1 type of metadata is requested.
   assert(Options.UAR || Options.Atomics);
 
   if (Options.UAR && !(FeatureMask & kSanitizerBinaryMetadataUAR)) {
     if (useAfterReturnUnsafe(I))
SanitizerBinaryMetadata::getSectionMarker(const Twine &MarkerName, Type *Ty) {
   auto *Marker = new GlobalVariable(Mod, Ty, /*isConstant=*/false,
                                     GlobalVariable::ExternalWeakLinkage,
                                     /*Initializer=*/nullptr, MarkerName);
   Marker->setVisibility(GlobalValue::HiddenVisibility);
   return Marker;
 }
 
 StringRef SanitizerBinaryMetadata::getSectionName(StringRef SectionSuffix) {
-  // FIXME: Other TargetTriple (req. string pool)
-  return SectionSuffix;
+  // FIXME: Other TargetTriples.
+  // Request ULEB128 encoding for all integer constants.
+  return StringPool.save(SectionSuffix + "!C");
 }
 
 Twine SanitizerBinaryMetadata::getSectionStart(StringRef SectionSuffix) {
   return "__start_" + SectionSuffix;
 }
 
 Twine SanitizerBinaryMetadata::getSectionEnd(StringRef SectionSuffix) {
   return "__stop_" + SectionSuffix;

llvm/test/CodeGen/X86/pcsections.ll

entry:
   store volatile i32 %inc, ptr @foo, align 4
   store atomic volatile i32 1, ptr @foo monotonic, align 4, !pcsections !0
   %1 = load volatile i32, ptr @foo, align 4
   %dec = add nsw i32 %1, -1
   store volatile i32 %dec, ptr @foo, align 4
   ret void
 }
 
+define void @multiple_uleb128() !pcsections !6 {
+; CHECK-LABEL: multiple_uleb128:
+; CHECK:       .section section_aux,"awo",@progbits,.text
+; CHECK-NEXT:  .Lpcsection_base8:
+; DEFCM-NEXT:  .long  .Lfunc_begin3-.Lpcsection_base8
+; LARGE-NEXT:  .quad  .Lfunc_begin3-.Lpcsection_base8
+; CHECK-NEXT:  .uleb128 .Lfunc_end6-.Lfunc_begin3
+; CHECK-NEXT:  .byte  42
+; CHECK-NEXT:  .ascii "\345\216&"
+; CHECK-NEXT:  .byte  255
+; CHECK-NEXT:  .section section_aux_21264,"awo",@progbits,.text
+; CHECK-NEXT:  .Lpcsection_base9:
+; DEFCM-NEXT:  .long  .Lfunc_begin3-.Lpcsection_base9
+; LARGE-NEXT:  .quad  .Lfunc_begin3-.Lpcsection_base9
+; CHECK-NEXT:  .long  .Lfunc_end6-.Lfunc_begin3
+; CHECK-NEXT:  .long  21264
+; CHECK-NEXT:  .text
+entry:
+  ret void
+}
+
 !0 = !{!"section_no_aux"}
 !1 = !{!"section_aux", !3}
 !2 = !{!"section_aux_42", !4, !"section_aux_21264", !5}
 !3 = !{i32 10, i32 20, i32 30}
 !4 = !{i32 42}
 !5 = !{i32 21264}
+!6 = !{!"section_aux!C", !7, !"section_aux_21264", !5}
+!7 = !{i64 42, i32 624485, i8 255}

llvm/test/Instrumentation/SanitizerBinaryMetadata/atomics.ll

 ; CHECK-DAG: entry:
 ; CHECK-NEXT:  br i1 icmp ne (ptr @__sanitizer_metadata_covered_del, ptr null), label %callfunc, label %ret
 ; CHECK-DAG: callfunc:
 ; CHECK-NEXT:  call void @__sanitizer_metadata_covered_del(i32 2, ptr @__start_sanmd_covered, ptr @__stop_sanmd_covered)
 ; CHECK-NEXT:  br label %ret
 ; CHECK-DAG: ret:
 ; CHECK-NEXT:  ret void
 
-; CHECK: !0 = !{!"sanmd_covered", !1}
-; CHECK: !1 = !{i8 1}
-; CHECK: !2 = !{!"sanmd_atomics"}
+; CHECK: !0 = !{!"sanmd_covered!C", !1}
+; CHECK: !1 = !{i64 1}
+; CHECK: !2 = !{!"sanmd_atomics!C"}