diff --git a/llvm/include/llvm/IR/Module.h b/llvm/include/llvm/IR/Module.h
--- a/llvm/include/llvm/IR/Module.h
+++ b/llvm/include/llvm/IR/Module.h
@@ -493,10 +493,12 @@
   void addModuleFlag(ModFlagBehavior Behavior, StringRef Key, Constant *Val);
   void addModuleFlag(ModFlagBehavior Behavior, StringRef Key, uint32_t Val);
   void addModuleFlag(MDNode *Node);
+  /// Like addModuleFlag but replaces the old module flag if it already exists.
+  void setModuleFlag(ModFlagBehavior Behavior, StringRef Key, Metadata *Val);
 
-/// @}
-/// @name Materialization
-/// @{
+  /// @}
+  /// @name Materialization
+  /// @{
 
   /// Sets the GVMaterializer to GVM. This module must not yet have a
   /// Materializer. To reset the materializer for a module that already has one,
diff --git a/llvm/lib/IR/Module.cpp b/llvm/lib/IR/Module.cpp
--- a/llvm/lib/IR/Module.cpp
+++ b/llvm/lib/IR/Module.cpp
@@ -358,6 +358,24 @@
   getOrInsertModuleFlagsMetadata()->addOperand(Node);
 }
 
+void Module::setModuleFlag(ModFlagBehavior Behavior, StringRef Key,
+                           Metadata *Val) {
+  NamedMDNode *ModFlags = getOrInsertModuleFlagsMetadata();
+  // Replace the flag if it already exists.
+  for (unsigned I = 0, E = ModFlags->getNumOperands(); I != E; ++I) {
+    MDNode *Flag = ModFlags->getOperand(I);
+    ModFlagBehavior MFB;
+    if (Flag->getNumOperands() == 3 &&
+        isValidModFlagBehavior(Flag->getOperand(0), MFB) && MFB == Behavior &&
+        dyn_cast_or_null<MDString>(Flag->getOperand(1)) &&
+        cast<MDString>(Flag->getOperand(1))->getString() == Key) {
+      Flag->replaceOperandWith(2, Val);
+      return;
+    }
+  }
+  addModuleFlag(Behavior, Key, Val);
+}
+
 void Module::setDataLayout(StringRef Desc) {
   DL.reset(Desc);
 }
@@ -547,9 +565,9 @@
 
 void Module::setProfileSummary(Metadata *M, ProfileSummary::Kind Kind) {
   if (Kind == ProfileSummary::PSK_CSInstr)
-    addModuleFlag(ModFlagBehavior::Error, "CSProfileSummary", M);
+    setModuleFlag(ModFlagBehavior::Error, "CSProfileSummary", M);
   else
-    addModuleFlag(ModFlagBehavior::Error, "ProfileSummary", M);
+    setModuleFlag(ModFlagBehavior::Error, "ProfileSummary", M);
 }
 
 Metadata *Module::getProfileSummary(bool IsCS) {
diff --git a/llvm/unittests/IR/ModuleTest.cpp b/llvm/unittests/IR/ModuleTest.cpp
--- a/llvm/unittests/IR/ModuleTest.cpp
+++ b/llvm/unittests/IR/ModuleTest.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/IR/Module.h"
+#include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/RandomNumberGenerator.h"
@@ -72,4 +73,52 @@
                          RandomStreams[1].begin()));
 }
 
+TEST(ModuleTest, setModuleFlag) {
+  LLVMContext Context;
+  Module M("M", Context);
+  StringRef Key = "Key";
+  Metadata *Val1 = MDString::get(Context, "Val1");
+  Metadata *Val2 = MDString::get(Context, "Val2");
+  EXPECT_EQ(nullptr, M.getModuleFlag(Key));
+  M.setModuleFlag(Module::ModFlagBehavior::Error, Key, Val1);
+  EXPECT_EQ(Val1, M.getModuleFlag(Key));
+  M.setModuleFlag(Module::ModFlagBehavior::Error, Key, Val2);
+  EXPECT_EQ(Val2, M.getModuleFlag(Key));
+}
+
+const char *IRString = R"IR(
+  !llvm.module.flags = !{!0}
+
+  !0 = !{i32 1, !"ProfileSummary", !1}
+  !1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
+  !2 = !{!"ProfileFormat", !"SampleProfile"}
+  !3 = !{!"TotalCount", i64 10000}
+  !4 = !{!"MaxCount", i64 10}
+  !5 = !{!"MaxInternalCount", i64 1}
+  !6 = !{!"MaxFunctionCount", i64 1000}
+  !7 = !{!"NumCounts", i64 200}
+  !8 = !{!"NumFunctions", i64 3}
+  !9 = !{!"DetailedSummary", !10}
+  !10 = !{!11, !12, !13}
+  !11 = !{i32 10000, i64 1000, i32 1}
+  !12 = !{i32 990000, i64 300, i32 10}
+  !13 = !{i32 999999, i64 5, i32 100}
+)IR";
+
+TEST(ModuleTest, setProfileSummary) {
+  SMDiagnostic Err;
+  LLVMContext Context;
+  std::unique_ptr<Module> M = parseAssemblyString(IRString, Err, Context);
+  auto *PS = ProfileSummary::getFromMD(M->getProfileSummary(/*IsCS*/ false));
+  EXPECT_NE(nullptr, PS);
+  EXPECT_EQ(false, PS->isPartialProfile());
+  PS->setPartialProfile(true);
+  M->setProfileSummary(PS->getMD(Context), ProfileSummary::PSK_Sample);
+  delete PS;
+  PS = ProfileSummary::getFromMD(M->getProfileSummary(/*IsCS*/ false));
+  EXPECT_NE(nullptr, PS);
+  EXPECT_EQ(true, PS->isPartialProfile());
+  delete PS;
+}
+
 } // end namespace