[IR] Make Module::setProfileSummary to replace an existing ProfileSummary flag.

[hjyamauchi]

Summary

Module::setProfileSummary currently calls addModuelFlag. This prevents from
updating the ProfileSummary metadata in the module and results in a second
ProfileSummary added instead of replacing an existing one. I don't think this is
the expected behavior. It prevents updating the ProfileSummary and it does not
make sense to have more than one. To address this, add Module::setModuleFlag and
use it from setProfileSummary.

[wmi]

llvm/lib/IR/Module.cpp

361–362

It has some code similar as getModuleFlagsMetadata. Can they be shared?

[davidxl]

Can CS FDO profile summary coexist with regular summary?

[hjyamauchi]

Address comments.

[hjyamauchi]

In D79902#2044401, @davidxl wrote:

Can CS FDO profile summary coexist with regular summary?

Yes, it uses a different key "CSProfileSummary". So it does not conflict with the regular summary.

[xur]

I noticed the same issue when I implemented CSFDO -- I initially reused the same summary name for CSFDO as in FDO.
Exiting of two module flags with the same name sometime breaks the build (not always, if I remember correctly).

This fix looks to me.

[davidxl]

lgtm

llvm/include/llvm/IR/Module.h

/// @{
   NamedMDNode *getOrInsertModuleFlagsMetadata();
 
   /// Add a module-level flag to the module-level flags metadata. It will create
   /// the module-level flags named metadata if it doesn't already exist.
   void addModuleFlag(ModFlagBehavior Behavior, StringRef Key, Metadata *Val);
   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
   /// @{
 
   /// Sets the GVMaterializer to GVM. This module must not yet have a
   /// Materializer. To reset the materializer for a module that already has one,
   /// call materializeAll first. Destroying this module will destroy
   /// its materializer without materializing any more GlobalValues. Without
   /// destroying the Module, there is no way to detach or destroy a materializer
   /// without materializing all the GVs it controls, to avoid leaving orphan
   /// unmaterialized GVs.

llvm/lib/IR/Module.cpp

void Module::addModuleFlag(MDNode *Node) {
   assert(Node->getNumOperands() == 3 &&
          "Invalid number of operands for module flag!");
   assert(mdconst::hasa<ConstantInt>(Node->getOperand(0)) &&
          isa<MDString>(Node->getOperand(1)) &&
          "Invalid operand types for module flag!");
   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);
 }
 
 void Module::setDataLayout(const DataLayout &Other) { DL = Other; }
 
 const DataLayout &Module::getDataLayout() const { return DL; }
 
void Module::setCodeModel(CodeModel::Model CL) {
   // because the compiler would have to generate additional code (to span
   // longer jumps) if a larger code model is used with a smaller one.
   // Therefore we will treat attempts to mix code models as an error.
   addModuleFlag(ModFlagBehavior::Error, "Code Model", CL);
 }
 
 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) {
   return (IsCS ? getModuleFlag("CSProfileSummary")
                : getModuleFlag("ProfileSummary"));
 }
 
 bool Module::getSemanticInterposition() const {

llvm/unittests/IR/ModuleTest.cpp

 //===- unittests/IR/ModuleTest.cpp - Module unit tests --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/IR/Module.h"
+#include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/RandomNumberGenerator.h"
 #include "gtest/gtest.h"
 
 #include <random>
 
 using namespace llvm;
for (auto &RandomStream : RandomStreams) {
     std::generate(RandomStream.begin(), RandomStream.end(),
                   [&]() { return dist(*RNG); });
   }
 
   EXPECT_TRUE(std::equal(RandomStreams[0].begin(), RandomStreams[0].end(),
                          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