1 files changed, 16 insertions, 0 deletions
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp b/llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp
index b461c3c4bfdc..f5e12fd960d0 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp
@@ -11,6 +11,7 @@
 #include "AMDGPUPerfHintAnalysis.h"
 #include "AMDGPUSubtarget.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/IR/Constants.h"
 #include "llvm/Target/TargetMachine.h"
 
 using namespace llvm;
@@ -101,6 +102,21 @@ void AMDGPUMachineFunction::allocateModuleLDSGlobal(const Function &F) {
   }
 }
 
+Optional<uint32_t>
+AMDGPUMachineFunction::getLDSKernelIdMetadata(const Function &F) {
+  auto MD = F.getMetadata("llvm.amdgcn.lds.kernel.id");
+  if (MD && MD->getNumOperands() == 1) {
+    ConstantInt *KnownSize = mdconst::extract<ConstantInt>(MD->getOperand(0));
+    if (KnownSize) {
+      uint64_t V = KnownSize->getZExtValue();
+      if (V <= UINT32_MAX) {
+        return V;
+      }
+    }
+  }
+  return {};
+}
+
 void AMDGPUMachineFunction::setDynLDSAlign(const DataLayout &DL,
                                            const GlobalVariable &GV) {
   assert(DL.getTypeAllocSize(GV.getValueType()).isZero());