[IR] Fix Module::setModuleFlag for uniqued metadata

[asavonic]

Module::setModuleFlag is supposed to change a single module. However, when an MDNode has the same value in more than one module in the same LLVMContext, such MDNode is shared (uniqued) across all of them. Therefore MDNode::replaceOperandWith changes all modules that share the same MDNode.

This used to cause problems for #86212, where a module is marked as "upgraded" via a module flag. When this flag is shared across multiple modules, all of them are marked, yet some may not have been processed at all.

After the patch we now construct a new MDNode and replace the old one.

[llvmbot]

@llvm/pr-subscribers-llvm-ir

Author: Andrew Savonichev (asavonic)

Changes

Module::setModuleFlag is supposed to change a single module. However, when an MDNode has the same value in more than one module in the same LLVMContext, such MDNode is shared (uniqued) across all of them. Therefore MDNode::replaceOperandWith changes all modules that share the same MDNode.

This used to cause problems for #86212, where a module is marked as "upgraded" via a module flag. When this flag is shared across multiple modules, all of them are marked, yet some may not have been processed at all.

After the patch we now construct a new MDNode and replace the old one.

---

Full diff: https://github.com/llvm/llvm-project/pull/164580.diff

2 Files Affected:

• (modified) llvm/lib/IR/Module.cpp (+7-2)
• (modified) llvm/unittests/IR/ModuleTest.cpp (+30)

diff --git a/llvm/lib/IR/Module.cpp b/llvm/lib/IR/Module.cpp
index 30b5e48652b28..e19336efd1c8a 100644
--- a/llvm/lib/IR/Module.cpp
+++ b/llvm/lib/IR/Module.cpp
@@ -403,9 +403,14 @@ void Module::setModuleFlag(ModFlagBehavior Behavior, StringRef Key,
                            Metadata *Val) {
   NamedMDNode *ModFlags = getOrInsertModuleFlagsMetadata();
   // Replace the flag if it already exists.
-  for (MDNode *Flag : ModFlags->operands()) {
+  for (unsigned i = 0; i < ModFlags->getNumOperands(); ++i) {
+    MDNode *Flag = ModFlags->getOperand(i);
     if (cast<MDString>(Flag->getOperand(1))->getString() == Key) {
-      Flag->replaceOperandWith(2, Val);
+      Type *Int32Ty = Type::getInt32Ty(Context);
+      Metadata *Ops[3] = {
+          ConstantAsMetadata::get(ConstantInt::get(Int32Ty, Behavior)),
+          MDString::get(Context, Key), Val};
+      ModFlags->setOperand(i, MDNode::get(Context, Ops));
       return;
     }
   }
diff --git a/llvm/unittests/IR/ModuleTest.cpp b/llvm/unittests/IR/ModuleTest.cpp
index 36c356730d27a..30eda738020d0 100644
--- a/llvm/unittests/IR/ModuleTest.cpp
+++ b/llvm/unittests/IR/ModuleTest.cpp
@@ -103,6 +103,36 @@ TEST(ModuleTest, setModuleFlagInt) {
   EXPECT_EQ(Val2, A2->getZExtValue());
 }
 
+TEST(ModuleTest, setModuleFlagTwoMod) {
+  LLVMContext Context;
+  Module MA("MA", Context);
+  Module MB("MB", Context);
+  StringRef Key = "Key";
+  uint32_t Val1 = 1;
+  uint32_t Val2 = 2;
+
+  // Set a flag to MA
+  EXPECT_EQ(nullptr, MA.getModuleFlag(Key));
+  MA.setModuleFlag(Module::ModFlagBehavior::Error, Key, Val1);
+  auto A1 = mdconst::extract_or_null<ConstantInt>(MA.getModuleFlag(Key));
+  EXPECT_EQ(Val1, A1->getZExtValue());
+
+  // Set a flag to MB
+  EXPECT_EQ(nullptr, MB.getModuleFlag(Key));
+  MB.setModuleFlag(Module::ModFlagBehavior::Error, Key, Val1);
+  auto B1 = mdconst::extract_or_null<ConstantInt>(MB.getModuleFlag(Key));
+  EXPECT_EQ(Val1, B1->getZExtValue());
+
+  // Change the flag of MA
+  MA.setModuleFlag(Module::ModFlagBehavior::Error, Key, Val2);
+  auto A2 = mdconst::extract_or_null<ConstantInt>(MA.getModuleFlag(Key));
+  EXPECT_EQ(Val2, A2->getZExtValue());
+
+  // MB should keep the original flag value
+  auto B2 = mdconst::extract_or_null<ConstantInt>(MB.getModuleFlag(Key));
+  EXPECT_EQ(Val1, B2->getZExtValue());
+}
+
 const char *IRString = R"IR(
   !llvm.module.flags = !{!0}
 

[DanielKristofKiss]

nice catch, I didn't expect cross module update by the setModuleFlag.
LGTM