[llvm] bd2430b - [IR] Allow type change in ValueAsMetadata::handleRAUW (#76969)

[via llvm-commits]

Author: Jannik Silvanus
Date: 2024-01-18T16:01:23+01:00
New Revision: bd2430b4216f4f4f4432bc86cc8d7e1b7a71f56c

URL: https://github.com/llvm/llvm-project/commit/bd2430b4216f4f4f4432bc86cc8d7e1b7a71f56c
DIFF: https://github.com/llvm/llvm-project/commit/bd2430b4216f4f4f4432bc86cc8d7e1b7a71f56c.diff

LOG: [IR] Allow type change in ValueAsMetadata::handleRAUW (#76969)

`ValueAsMetadata::handleRAUW` is a mechanism to replace all metadata
referring to one value by a different value.

Relax an assert that used to enforce the old and new value to have the
same type.
This seems to be a sanity plausibility assert only, as the
implementation actually supports mismatching types.

This is motivated by a downstream mechanism where we use poison
ValueAsMetadata values to annotate pointee types of opaque pointer
function arguments.

When replacing one type with a different one to work around DXIL vs LLVM
incompatibilities, we need to update type annotations, and handleRAUW is
more efficient than creating new MD nodes.

Added: 
    

Modified: 
    llvm/lib/IR/Metadata.cpp
    llvm/unittests/IR/MetadataTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
index a9f64f5442341c..eff4fd371a73ff 100644
--- a/llvm/lib/IR/Metadata.cpp
+++ b/llvm/lib/IR/Metadata.cpp
@@ -530,7 +530,7 @@ void ValueAsMetadata::handleRAUW(Value *From, Value *To) {
   assert(From && "Expected valid value");
   assert(To && "Expected valid value");
   assert(From != To && "Expected changed value");
-  assert(From->getType() == To->getType() && "Unexpected type change");
+  assert(&From->getContext() == &To->getContext() && "Expected same context");
 
   LLVMContext &Context = From->getType()->getContext();
   auto &Store = Context.pImpl->ValuesAsMetadata;

diff  --git a/llvm/unittests/IR/MetadataTest.cpp b/llvm/unittests/IR/MetadataTest.cpp
index 1570631287b2a7..767dd1a59d2b9e 100644
--- a/llvm/unittests/IR/MetadataTest.cpp
+++ b/llvm/unittests/IR/MetadataTest.cpp
@@ -3724,6 +3724,25 @@ TEST_F(ValueAsMetadataTest, UpdatesOnRAUW) {
   EXPECT_TRUE(MD->getValue() == GV1.get());
 }
 
+TEST_F(ValueAsMetadataTest, handleRAUWWithTypeChange) {
+  // Test that handleRAUW supports type changes.
+  // This is helpful in cases where poison values are used to encode
+  // types in metadata, e.g. in type annotations.
+  // Changing the type stored in metadata requires to change the type of
+  // the stored poison value.
+  auto *I32Poison = PoisonValue::get(Type::getInt32Ty(Context));
+  auto *I64Poison = PoisonValue::get(Type::getInt64Ty(Context));
+  auto *MD = ConstantAsMetadata::get(I32Poison);
+
+  EXPECT_EQ(MD->getValue(), I32Poison);
+  EXPECT_NE(MD->getValue(), I64Poison);
+
+  ValueAsMetadata::handleRAUW(I32Poison, I64Poison);
+
+  EXPECT_NE(MD->getValue(), I32Poison);
+  EXPECT_EQ(MD->getValue(), I64Poison);
+}
+
 TEST_F(ValueAsMetadataTest, TempTempReplacement) {
   // Create a constant.
   ConstantAsMetadata *CI =