2 files changed, 12 insertions, 6 deletions
diff --git a/llvm/lib/IR/Module.cpp b/llvm/lib/IR/Module.cpp
index b6dd7ab70a8..dbf4035ac7c 100644
--- a/llvm/lib/IR/Module.cpp
+++ b/llvm/lib/IR/Module.cpp
@@ -531,12 +531,16 @@ void Module::setCodeModel(CodeModel::Model CL) {
   addModuleFlag(ModFlagBehavior::Error, "Code Model", CL);
 }
 
-void Module::setProfileSummary(Metadata *M) {
-  addModuleFlag(ModFlagBehavior::Error, "ProfileSummary", M);
+void Module::setProfileSummary(Metadata *M, ProfileSummary::Kind Kind) {
+  if (Kind == ProfileSummary::PSK_CSInstr)
+    addModuleFlag(ModFlagBehavior::Error, "CSProfileSummary", M);
+  else
+    addModuleFlag(ModFlagBehavior::Error, "ProfileSummary", M);
 }
 
-Metadata *Module::getProfileSummary() {
-  return getModuleFlag("ProfileSummary");
+Metadata *Module::getProfileSummary(bool IsCS) {
+  return (IsCS ? getModuleFlag("CSProfileSummary")
+               : getModuleFlag("ProfileSummary"));
 }
 
 void Module::setOwnedMemoryBuffer(std::unique_ptr<MemoryBuffer> MB) {
diff --git a/llvm/lib/IR/ProfileSummary.cpp b/llvm/lib/IR/ProfileSummary.cpp
index d8596924719..11d95ac19be 100644
--- a/llvm/lib/IR/ProfileSummary.cpp
+++ b/llvm/lib/IR/ProfileSummary.cpp
@@ -21,8 +21,6 @@
 
 using namespace llvm;
 
-const char *ProfileSummary::KindStr[2] = {"InstrProf", "SampleProfile"};
-
 // Return an MDTuple with two elements. The first element is a string Key and
 // the second is a uint64_t Value.
 static Metadata *getKeyValMD(LLVMContext &Context, const char *Key,
@@ -68,6 +66,7 @@ Metadata *ProfileSummary::getDetailedSummaryMD(LLVMContext &Context) {
 // "SampleProfile"). The rest of the elements of the outer MDTuple are specific
 // to the kind of profile summary as returned by getFormatSpecificMD.
 Metadata *ProfileSummary::getMD(LLVMContext &Context) {
+  const char *KindStr[3] = {"InstrProf", "CSInstrProf", "SampleProfile"};
   Metadata *Components[] = {
     getKeyValMD(Context, "ProfileFormat", KindStr[PSK]),
     getKeyValMD(Context, "TotalCount", getTotalCount()),
@@ -153,6 +152,9 @@ ProfileSummary *ProfileSummary::getFromMD(Metadata *MD) {
   else if (isKeyValuePair(dyn_cast_or_null<MDTuple>(FormatMD), "ProfileFormat",
                           "InstrProf"))
     SummaryKind = PSK_Instr;
+  else if (isKeyValuePair(dyn_cast_or_null<MDTuple>(FormatMD), "ProfileFormat",
+                          "CSInstrProf"))
+    SummaryKind = PSK_CSInstr;
   else
     return nullptr;