[JITLink] Track section alignment and make sure it is respected during layout.

Previously we had only honored alignments on individual atoms, but
tools/runtimes may assume that the section alignment is respected too.

llvm-svn: 360555

4 files changed, 43 insertions, 11 deletions

diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
index b3477f1ee4c..96e074da122 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
@@ -241,6 +241,10 @@ Error JITLinkerBase::allocateSegments(const SegmentLayoutMap &Layout) {
     for (auto &SI : SegLayout.ContentSections) {
       assert(!SI.S->atoms_empty() && "Sections in layout must not be empty");
       assert(!SI.Atoms.empty() && "Section layouts must not be empty");
+
+      // Bump to section alignment before processing atoms.
+      SegContentSize = alignTo(SegContentSize, SI.S->getAlignment());
+
       for (auto *DA : SI.Atoms) {
         SegContentSize = alignTo(SegContentSize, DA->getAlignment());
         SegContentSize += DA->getSize();
@@ -249,15 +253,22 @@ Error JITLinkerBase::allocateSegments(const SegmentLayoutMap &Layout) {
 
     // Get segment content alignment.
     unsigned SegContentAlign = 1;
-    if (!SegLayout.ContentSections.empty())
+    if (!SegLayout.ContentSections.empty()) {
+      auto &FirstContentSection = SegLayout.ContentSections.front();
       SegContentAlign =
-          SegLayout.ContentSections.front().Atoms.front()->getAlignment();
+          std::max(FirstContentSection.S->getAlignment(),
+                   FirstContentSection.Atoms.front()->getAlignment());
+    }
 
     // Calculate segment zero-fill size.
     uint64_t SegZeroFillSize = 0;
     for (auto &SI : SegLayout.ZeroFillSections) {
       assert(!SI.S->atoms_empty() && "Sections in layout must not be empty");
       assert(!SI.Atoms.empty() && "Section layouts must not be empty");
+
+      // Bump to section alignment before processing atoms.
+      SegZeroFillSize = alignTo(SegZeroFillSize, SI.S->getAlignment());
+
       for (auto *DA : SI.Atoms) {
         SegZeroFillSize = alignTo(SegZeroFillSize, DA->getAlignment());
         SegZeroFillSize += DA->getSize();
@@ -266,9 +277,13 @@ Error JITLinkerBase::allocateSegments(const SegmentLayoutMap &Layout) {
 
     // Calculate segment zero-fill alignment.
     uint32_t SegZeroFillAlign = 1;
-    if (!SegLayout.ZeroFillSections.empty())
+
+    if (!SegLayout.ZeroFillSections.empty()) {
+      auto &FirstZeroFillSection = SegLayout.ZeroFillSections.front();
       SegZeroFillAlign =
-          SegLayout.ZeroFillSections.front().Atoms.front()->getAlignment();
+          std::max(FirstZeroFillSection.S->getAlignment(),
+                   FirstZeroFillSection.Atoms.front()->getAlignment());
+    }
 
     if (SegContentSize == 0)
       SegContentAlign = SegZeroFillAlign;
@@ -314,12 +329,14 @@ Error JITLinkerBase::allocateSegments(const SegmentLayoutMap &Layout) {
         Alloc->getTargetMemory(static_cast<sys::Memory::ProtectionFlags>(Prot));
 
     for (auto *SIList : {&SL.ContentSections, &SL.ZeroFillSections})
-      for (auto &SI : *SIList)
+      for (auto &SI : *SIList) {
+        AtomTargetAddr = alignTo(AtomTargetAddr, SI.S->getAlignment());
         for (auto *DA : SI.Atoms) {
           AtomTargetAddr = alignTo(AtomTargetAddr, DA->getAlignment());
           DA->setAddress(AtomTargetAddr);
           AtomTargetAddr += DA->getSize();
         }
+      }
   }
 
   return Error::success();
diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.h b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.h
index 9d8238cf2e5..e6fd6e38f7a 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.h
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.h
@@ -150,7 +150,7 @@ private:
       auto SegMem = Alloc.getWorkingMemory(
           static_cast<sys::Memory::ProtectionFlags>(Prot));
       char *LastAtomEnd = SegMem.data();
-      char *AtomDataPtr = nullptr;
+      char *AtomDataPtr = LastAtomEnd;
 
       LLVM_DEBUG({
         dbgs() << "  Processing segment "
@@ -162,8 +162,16 @@ private:
 
       for (auto &SI : SegLayout.ContentSections) {
         LLVM_DEBUG(dbgs() << "    " << SI.S->getName() << ":\n");
+
+        AtomDataPtr += alignmentAdjustment(AtomDataPtr, SI.S->getAlignment());
+
+        LLVM_DEBUG({
+          dbgs() << "      Bumped atom pointer to " << (const void *)AtomDataPtr
+                 << " to meet section alignment "
+                 << " of " << SI.S->getAlignment() << "\n";
+        });
+
         for (auto *DA : SI.Atoms) {
-          AtomDataPtr = LastAtomEnd;
 
           // Align.
           AtomDataPtr += alignmentAdjustment(AtomDataPtr, DA->getAlignment());
@@ -209,6 +217,7 @@ private:
 
           // Update atom end pointer.
           LastAtomEnd = AtomDataPtr + DA->getContent().size();
+          AtomDataPtr = LastAtomEnd;
         }
       }
 
diff --git a/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.cpp b/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.cpp
index 0c19f8fee0d..b24b2256826 100644
--- a/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.cpp
@@ -85,7 +85,7 @@ MachOAtomGraphBuilder::MachOSection &MachOAtomGraphBuilder::getCommonSection() {
   if (!CommonSymbolsSection) {
     auto Prot = static_cast<sys::Memory::ProtectionFlags>(
         sys::Memory::MF_READ | sys::Memory::MF_WRITE);
-    auto &GenericSection = G->createSection("<common>", Prot, true);
+    auto &GenericSection = G->createSection("<common>", 1, Prot, true);
     CommonSymbolsSection = MachOSection(GenericSection);
   }
   return *CommonSymbolsSection;
@@ -102,6 +102,12 @@ Error MachOAtomGraphBuilder::parseSections() {
 
     unsigned SectionIndex = SecRef.getIndex() + 1;
 
+    uint32_t Align = SecRef.getAlignment();
+    if (!isPowerOf2_32(Align))
+      return make_error<JITLinkError>("Section " + Name +
+                                      " has non-power-of-2 "
+                                      "alignment");
+
     // FIXME: Get real section permissions
     // How, exactly, on MachO?
     sys::Memory::ProtectionFlags Prot;
@@ -112,7 +118,7 @@ Error MachOAtomGraphBuilder::parseSections() {
       Prot = static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
                                                        sys::Memory::MF_WRITE);
 
-    auto &GenericSection = G->createSection(Name, Prot, SecRef.isBSS());
+    auto &GenericSection = G->createSection(Name, Align, Prot, SecRef.isBSS());
 
     LLVM_DEBUG({
       dbgs() << "Adding section " << Name << ": "
diff --git a/llvm/lib/ExecutionEngine/JITLink/MachO_x86_64.cpp b/llvm/lib/ExecutionEngine/JITLink/MachO_x86_64.cpp
index 7b4ddc3019a..4010678c6d3 100644
--- a/llvm/lib/ExecutionEngine/JITLink/MachO_x86_64.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/MachO_x86_64.cpp
@@ -399,7 +399,7 @@ public:
 private:
   Section &getGOTSection() {
     if (!GOTSection)
-      GOTSection = &G.createSection("$__GOT", sys::Memory::MF_READ, false);
+      GOTSection = &G.createSection("$__GOT", 8, sys::Memory::MF_READ, false);
     return *GOTSection;
   }
 
@@ -407,7 +407,7 @@ private:
     if (!StubsSection) {
       auto StubsProt = static_cast<sys::Memory::ProtectionFlags>(
           sys::Memory::MF_READ | sys::Memory::MF_EXEC);
-      StubsSection = &G.createSection("$__STUBS", StubsProt, false);
+      StubsSection = &G.createSection("$__STUBS", 8, StubsProt, false);
     }
     return *StubsSection;
   }