[llvm] [MTE] generalize overalignment / size of MTE globals (PR #121957)

[Florian Mayer via llvm-commits]

https://github.com/fmayer created https://github.com/llvm/llvm-project/pull/121957

This was requested in https://github.com/llvm/llvm-project/pull/111918


>From 2feb85c15f64546cb6874e1ca0a1310bd1e1bedd Mon Sep 17 00:00:00 2001
From: Florian Mayer <fmayer at google.com>
Date: Tue, 7 Jan 2025 07:57:09 -0800
Subject: [PATCH] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20initia?=
 =?UTF-8?q?l=20version?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Created using spr 1.3.4
---
 llvm/include/llvm/IR/GlobalVariable.h      |  7 +++++
 llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp | 34 ++++++++++++++--------
 llvm/lib/IR/Globals.cpp                    | 10 +++++++
 3 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/llvm/include/llvm/IR/GlobalVariable.h b/llvm/include/llvm/IR/GlobalVariable.h
index 83e484816d7d4c..c1f8cdc4fd0a65 100644
--- a/llvm/include/llvm/IR/GlobalVariable.h
+++ b/llvm/include/llvm/IR/GlobalVariable.h
@@ -267,6 +267,13 @@ class GlobalVariable : public GlobalObject, public ilist_node<GlobalVariable> {
            getAttributes().hasAttribute("rodata-section");
   }
 
+  MaybeAlign getRequiredGlobalAlignment() {
+    return isTagged() && getAlign().valueOrOne() < 16 ? MaybeAlign(16)
+                                                      : std::nullopt;
+  }
+
+  std::optional<uint64_t> getRequiredGlobalSize();
+
   /// Get the custom code model raw value of this global.
   ///
   unsigned getCodeModelRaw() const {
diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index 7bd3fb33b47d2b..8e594a8f92b204 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -2398,17 +2398,23 @@ void AsmPrinter::emitRemarksSection(remarks::RemarkStreamer &RS) {
   OutStreamer->emitBinaryData(Buf);
 }
 
-static void tagGlobalDefinition(Module &M, GlobalVariable *G) {
-  Constant *Initializer = G->getInitializer();
+static uint64_t globalSize(const GlobalVariable &G) {
+  const Constant *Initializer = G.getInitializer();
   uint64_t SizeInBytes =
-      M.getDataLayout().getTypeAllocSize(Initializer->getType());
+      G.getParent()->getDataLayout().getTypeAllocSize(Initializer->getType());
 
-  uint64_t NewSize = alignTo(SizeInBytes, 16);
-  if (SizeInBytes != NewSize) {
+  return SizeInBytes;
+}
+
+static void tagGlobalDefinition(GlobalVariable *G) {
+  Module &M = *G->getParent();
+  uint64_t SizeInBytes = globalSize(*G);
+  if (auto NewSize = G->getRequiredGlobalSize()) {
+    assert(*NewSize > SizeInBytes);
     // Pad the initializer out to the next multiple of 16 bytes.
-    llvm::SmallVector<uint8_t> Init(NewSize - SizeInBytes, 0);
+    llvm::SmallVector<uint8_t> Init(*NewSize - SizeInBytes, 0);
     Constant *Padding = ConstantDataArray::get(M.getContext(), Init);
-    Initializer = ConstantStruct::getAnon({Initializer, Padding});
+    auto *Initializer = ConstantStruct::getAnon({G->getInitializer(), Padding});
     auto *NewGV = new GlobalVariable(
         M, Initializer->getType(), G->isConstant(), G->getLinkage(),
         Initializer, "", G, G->getThreadLocalMode(), G->getAddressSpace());
@@ -2422,8 +2428,10 @@ static void tagGlobalDefinition(Module &M, GlobalVariable *G) {
     G = NewGV;
   }
 
-  if (G->getAlign().valueOrOne() < 16)
-    G->setAlignment(Align(16));
+  if (auto Align = G->getRequiredGlobalAlignment()) {
+    assert(*Align > G->getAlign().valueOrOne());
+    G->setAlignment(*Align);
+  }
 
   // Ensure that tagged globals don't get merged by ICF - as they should have
   // different tags at runtime.
@@ -2438,12 +2446,14 @@ bool AsmPrinter::doFinalization(Module &M) {
 
   std::vector<GlobalVariable *> GlobalsToTag;
   for (GlobalVariable &G : M.globals()) {
-    if (G.isDeclaration() || !G.isTagged())
+    if (G.isDeclaration())
       continue;
-    GlobalsToTag.push_back(&G);
+    if (G.getRequiredGlobalAlignment().has_value() ||
+        G.getRequiredGlobalSize().has_value())
+      GlobalsToTag.push_back(&G);
   }
   for (GlobalVariable *G : GlobalsToTag)
-    tagGlobalDefinition(M, G);
+    tagGlobalDefinition(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
diff --git a/llvm/lib/IR/Globals.cpp b/llvm/lib/IR/Globals.cpp
index db5e1cb57b1bab..79ff120c250e40 100644
--- a/llvm/lib/IR/Globals.cpp
+++ b/llvm/lib/IR/Globals.cpp
@@ -540,6 +540,16 @@ void GlobalVariable::setCodeModel(CodeModel::Model CM) {
   assert(getCodeModel() == CM && "Code model representation error!");
 }
 
+std::optional<uint64_t> GlobalVariable::getRequiredGlobalSize() {
+  if (!isTagged())
+    return std::nullopt;
+  Constant *Initializer = getInitializer();
+  uint64_t SizeInBytes =
+      getParent()->getDataLayout().getTypeAllocSize(Initializer->getType());
+  uint64_t Aligned = alignTo(SizeInBytes, 16);
+  return Aligned != SizeInBytes ? std::optional(Aligned) : std::nullopt;
+}
+
 //===----------------------------------------------------------------------===//
 // GlobalAlias Implementation
 //===----------------------------------------------------------------------===//