[MTE] generalize overalignment / size of MTE globals

clang/test/CodeGen/memtag-globals-asm.cpp

@@ -53,8 +53,8 @@
 // CHECK-A: .memtag global_int
 // CHECK-A: .globl global_int
 // CHECK-A: .p2align 4, 0x0
-// CHECK-A: .zero 16
 // CHECK-A: .size global_int, 16
+// CHECK-A: .p2align 4, 0x0
 int global_int;
 // CHECK-B: .memtag _ZL9local_int
 // CHECK-B: .local _ZL9local_int
@@ -69,36 +69,37 @@ static char local_buffer[16];
 // CHECK-D: .p2align 4, 0x0
 // CHECK-D: _ZL22local_buffer_local_end:
 // CHECK-D: .xword _ZL12local_buffer+16
-// CHECK-D: .zero 8
 // CHECK-D: .size _ZL22local_buffer_local_end, 16
+// CHECK-D: .p2align 4, 0x0
 static char* local_buffer_local_end = &local_buffer[16];
 // CHECK-E: .memtag local_buffer_global_end
 // CHECK-E: .globl local_buffer_global_end
 // CHECK-E  .p2align 4, 0x0
 // CHECK-E: local_buffer_global_end:
 // CHECK-E: .xword _ZL12local_buffer+16
-// CHECK-E: .zero 8
 // CHECK-E: .size local_buffer_global_end, 16
+// CHECK-E: .p2align 4, 0x0
 char* local_buffer_global_end = &local_buffer[16];
 
 // CHECK-F: .memtag global_buffer
 // CHECK-F: .globl global_buffer
 // CHECK-F: .p2align 4, 0x0
-// CHECK-F: .zero 16
 // CHECK-F: .size global_buffer, 16
+// CHECK-F: .p2align 4, 0x0
 char global_buffer[16];
 // CHECK-G: .memtag _ZL23global_buffer_local_end
 // CHECK-G: .p2align 4, 0x0
 // CHECK-G: _ZL23global_buffer_local_end:
 // CHECK-G: .xword global_buffer+16
-// CHECK-G: .zero 8
 // CHECK-G: .size _ZL23global_buffer_local_end, 16
+// CHECK-G: .p2align 4, 0x0
 static char* global_buffer_local_end = &global_buffer[16];
 // CHECK-H: .memtag global_buffer_global_end
 // CHECK-H: .p2align 4, 0x0
 // CHECK-H: global_buffer_global_end:
 // CHECK-H: .xword global_buffer+16
 // CHECK-H: .size global_buffer_global_end, 16
+// CHECK-H: .p2align 4, 0x0
 char* global_buffer_global_end = &global_buffer[16];
 
 // CHECK-S-NOT: .memtag zero_sized
@@ -115,37 +116,37 @@ class MyClass {
 // CHECK-I: .memtag _ZN7MyClass12my_class_intE
 // CHECK-I: .globl _ZN7MyClass12my_class_intE
 // CHECK-I: .p2align 4, 0x0
-// CHECK-I: .zero 16
 // CHECK-I: .size _ZN7MyClass12my_class_intE, 16
+// CHECK-I: .p2align 4, 0x0
 int MyClass::my_class_int;
 // CHECK-NOT: .memtag _ZN7MyClass18my_class_const_intE
 const int MyClass::my_class_const_int = 1;
 
 // CHECK-J: .memtag global_my_class
 // CHECK-J: .globl global_my_class
 // CHECK-J: .p2align 4, 0x0
-// CHECK-J: .zero 8
 // CHECK-J: .size global_my_class, 16
+// CHECK-J: .p2align 4, 0x0
 MyClass global_my_class;
 // CHECK-K: .memtag _ZL14local_my_class
 // CHECK-K: .p2align 4, 0x0
-// CHECK-K: .zero 8
 // CHECK-K: .size _ZL14local_my_class, 16
+// CHECK-K: .p2align 4, 0x0
 static MyClass local_my_class;
 
 // CHECK-NOT: .memtag _ZL18local_const_string
 static const char local_const_string[] = "this is a local string";
 // CHECK-L: .memtag _ZL12local_string
 // CHECK-L: .p2align 4, 0x0
-// CHECK-L: .zero 9
 // CHECK-L: .size _ZL12local_string, 32
+// CHECK-L: .p2align 4, 0x0
 static char local_string[] = "this is a local string";
 
 // CHECK-M: .memtag global_atomic_int
 // CHECK-M: .globl global_atomic_int
 // CHECK-M: .p2align 4, 0x0
-// CHECK-M: .zero 16
 // CHECK-M: .size global_atomic_int, 16
+// CHECK-M: .p2align 4, 0x0
 _Atomic(int) global_atomic_int;
 // CHECK-N: .memtag _ZL16local_atomic_int
 // CHECK-N: .local _ZL16local_atomic_int
@@ -160,8 +161,8 @@ union MyUnion {
 // CHECK-O: .memtag global_union
 // CHECK-O: .globl global_union
 // CHECK-O: .p2align 4, 0x0
-// CHECK-O: .zero 16
 // CHECK-O: .size global_union, 16
+// CHECK-O: .p2align 4, 0x0
 MyUnion global_union;
 // CHECK-P: .memtag _ZL11local_union
 // CHECK-P: .local _ZL11local_union
@@ -277,6 +278,13 @@ int f(int x) {
       function_int;
 }
 
+// CHECK-S: .memtag global_overaligned_char
+// CHECK-S: .globl global_overaligned_char
+// CHECK-S: .p2align 4, 0x0
+// CHECK-S: .size global_overaligned_char, 16
+// CHECK-S: .p2align 4, 0x0
+alignas(16) char global_overaligned_char;
+
 typedef void (*func_t)(void);
 #define CONSTRUCTOR(section_name) \
   __attribute__((used)) __attribute__((section(section_name)))

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -879,6 +879,8 @@ class AsmPrinter : public MachineFunctionPass {
 
   bool DwarfUsesRelocationsAcrossSections = false;
 
+  void emitGlobalVariable(const GlobalVariable *GV, MaybeAlign OverAlignment);
+
   /// This method emits the header for the current function.
   virtual void emitFunctionHeader();
 
@@ -934,6 +936,9 @@ class AsmPrinter : public MachineFunctionPass {
   virtual bool shouldEmitWeakSwiftAsyncExtendedFramePointerFlags() const {
     return false;
   }
+  virtual MaybeAlign getRequiredGlobalAlignment(const GlobalVariable &GV) {
+    return std::nullopt;
+  };
 };
 
 } // end namespace llvm

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -715,8 +715,16 @@ MCSymbol *AsmPrinter::getSymbolPreferLocal(const GlobalValue &GV) const {
   return TM.getSymbol(&GV);
 }
 
-/// EmitGlobalVariable - Emit the specified global variable to the .s file.
 void AsmPrinter::emitGlobalVariable(const GlobalVariable *GV) {
+  MaybeAlign RequiredAlignment = getRequiredGlobalAlignment(*GV);
+  emitGlobalVariable(GV, RequiredAlignment);
+  if (RequiredAlignment)
+    OutStreamer->emitValueToAlignment(*RequiredAlignment);
+}
+
+/// EmitGlobalVariable - Emit the specified global variable to the .s file.
+void AsmPrinter::emitGlobalVariable(const GlobalVariable *GV,
+                                    MaybeAlign OverAlignment) {
   bool IsEmuTLSVar = TM.useEmulatedTLS() && GV->isThreadLocal();
   assert(!(IsEmuTLSVar && GV->hasCommonLinkage()) &&
          "No emulated TLS variables in the common section");
@@ -778,11 +786,20 @@ void AsmPrinter::emitGlobalVariable(const GlobalVariable *GV) {
 
   const DataLayout &DL = GV->getDataLayout();
   uint64_t Size = DL.getTypeAllocSize(GV->getValueType());
+  // In general, we *must* obey the specified alignment it.  Overaligning a
+  // global with a specified alignment is a prompt way to break globals
+  // emitted to sections and expected to be contiguous (e.g. ObjC metadata).
+  //
+  // If we get passed in an explicit overalignment, it is up to the caller
+  // to ensure that is not the case (i.e. that the GV is not in a section).
+  Align Alignment = getGVAlignment(GV, DL);
 
-  // If the alignment is specified, we *must* obey it.  Overaligning a global
-  // with a specified alignment is a prompt way to break globals emitted to
-  // sections and expected to be contiguous (e.g. ObjC metadata).
-  const Align Alignment = getGVAlignment(GV, DL);
+  if (OverAlignment) {
+    assert(!GV->hasSection());
+    Size = alignTo(Size, *OverAlignment);
+    if (Alignment < *OverAlignment)
+      Alignment = *OverAlignment;
+  }
 
   for (auto &Handler : Handlers)
     Handler->setSymbolSize(GVSym, Size);
@@ -2439,37 +2456,6 @@ static bool shouldTagGlobal(const llvm::GlobalVariable &G) {
   return globalSize(G) > 0;
 }
 
-static void tagGlobalDefinition(Module &M, GlobalVariable *G) {
-  uint64_t SizeInBytes = globalSize(*G);
-
-  uint64_t NewSize = alignTo(SizeInBytes, 16);
-  if (SizeInBytes != NewSize) {
-    // Pad the initializer out to the next multiple of 16 bytes.
-    llvm::SmallVector<uint8_t> Init(NewSize - SizeInBytes, 0);
-    Constant *Padding = ConstantDataArray::get(M.getContext(), Init);
-    Constant *Initializer = G->getInitializer();
-    Initializer = ConstantStruct::getAnon({Initializer, Padding});
-    auto *NewGV = new GlobalVariable(
-        M, Initializer->getType(), G->isConstant(), G->getLinkage(),
-        Initializer, "", G, G->getThreadLocalMode(), G->getAddressSpace());
-    NewGV->copyAttributesFrom(G);
-    NewGV->setComdat(G->getComdat());
-    NewGV->copyMetadata(G, 0);
-
-    NewGV->takeName(G);
-    G->replaceAllUsesWith(NewGV);
-    G->eraseFromParent();
-    G = NewGV;
-  }
-
-  if (G->getAlign().valueOrOne() < 16)
-    G->setAlignment(Align(16));
-
-  // Ensure that tagged globals don't get merged by ICF - as they should have
-  // different tags at runtime.
-  G->setUnnamedAddr(GlobalValue::UnnamedAddr::None);
-}
-
 static void removeMemtagFromGlobal(GlobalVariable &G) {
   auto Meta = G.getSanitizerMetadata();
   Meta.Memtag = false;
@@ -2482,7 +2468,6 @@ bool AsmPrinter::doFinalization(Module &M) {
   // we can conditionalize accesses based on whether or not it is nullptr.
   MF = nullptr;
 
-  std::vector<GlobalVariable *> GlobalsToTag;
   for (GlobalVariable &G : M.globals()) {
     if (G.isDeclaration() || !G.isTagged())
       continue;
@@ -2492,10 +2477,10 @@ bool AsmPrinter::doFinalization(Module &M) {
       assert(!G.isTagged());
       continue;
     }
-    GlobalsToTag.push_back(&G);
+    // Ensure that tagged globals don't get merged by ICF - as they should have
+    // different tags at runtime.
+    G.setUnnamedAddr(GlobalValue::UnnamedAddr::None);
   }
-  for (GlobalVariable *G : GlobalsToTag)
-    tagGlobalDefinition(M, G);
 
   // Gather all GOT equivalent globals in the module. We really need two
   // passes over the globals: one to compute and another to avoid its emission

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -282,6 +282,7 @@ class AArch64AsmPrinter : public AsmPrinter {
 
   void emitFunctionBodyEnd() override;
   void emitGlobalAlias(const Module &M, const GlobalAlias &GA) override;
+  MaybeAlign getRequiredGlobalAlignment(const GlobalVariable &GV) override;
 
   MCSymbol *GetCPISymbol(unsigned CPID) const override;
   void emitEndOfAsmFile(Module &M) override;
@@ -1463,6 +1464,11 @@ void AArch64AsmPrinter::emitGlobalAlias(const Module &M,
   AsmPrinter::emitGlobalAlias(M, GA);
 }
 
+MaybeAlign
+AArch64AsmPrinter::getRequiredGlobalAlignment(const GlobalVariable &GV) {
+  return GV.isTagged() ? MaybeAlign(16) : std::nullopt;
+}
+
 /// Small jump tables contain an unsigned byte or half, representing the offset
 /// from the lowest-addressed possible destination to the desired basic
 /// block. Since all instructions are 4-byte aligned, this is further compressed