1 files changed, 32 insertions, 0 deletions
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUHSATargetObjectFile.cpp b/llvm/lib/Target/AMDGPU/AMDGPUHSATargetObjectFile.cpp
index fa54f4a017c..ee42eaacf54 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUHSATargetObjectFile.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUHSATargetObjectFile.cpp
@@ -8,6 +8,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPUHSATargetObjectFile.h"
+#include "AMDGPU.h"
 #include "Utils/AMDGPUBaseInfo.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCSectionELF.h"
@@ -22,6 +23,29 @@ void AMDGPUHSATargetObjectFile::Initialize(MCContext &Ctx,
 
   TextSection = AMDGPU::getHSATextSection(Ctx);
 
+  DataGlobalAgentSection = AMDGPU::getHSADataGlobalAgentSection(Ctx);
+  DataGlobalProgramSection = AMDGPU::getHSADataGlobalProgramSection(Ctx);
+
+}
+
+bool AMDGPUHSATargetObjectFile::isAgentAllocationSection(
+    const char *SectionName) const {
+  return cast<MCSectionELF>(DataGlobalAgentSection)
+      ->getSectionName()
+      .equals(SectionName);
+}
+
+bool AMDGPUHSATargetObjectFile::isAgentAllocation(const GlobalValue *GV) const {
+  // Read-only segments can only have agent allocation.
+  return AMDGPU::isReadOnlySegment(GV) ||
+         (AMDGPU::isGlobalSegment(GV) && GV->hasSection() &&
+          isAgentAllocationSection(GV->getSection()));
+}
+
+bool AMDGPUHSATargetObjectFile::isProgramAllocation(
+    const GlobalValue *GV) const {
+  // The default for global segments is program allocation.
+  return AMDGPU::isGlobalSegment(GV) && !isAgentAllocation(GV);
 }
 
 MCSection *AMDGPUHSATargetObjectFile::SelectSectionForGlobal(
@@ -31,5 +55,13 @@ MCSection *AMDGPUHSATargetObjectFile::SelectSectionForGlobal(
   if (Kind.isText() && !GV->hasComdat())
     return getTextSection();
 
+  if (AMDGPU::isGlobalSegment(GV)) {
+    if (isAgentAllocation(GV))
+      return DataGlobalAgentSection;
+
+    if (isProgramAllocation(GV))
+      return DataGlobalProgramSection;
+  }
+
   return TargetLoweringObjectFileELF::SelectSectionForGlobal(GV, Kind, Mang, TM);
 }