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

[Jannik Silvanus via llvm-commits]

https://github.com/jasilvanus updated https://github.com/llvm/llvm-project/pull/76969

>From 100133e19932cf083b8a55f938bbb813b57fa28a Mon Sep 17 00:00:00 2001
From: Jannik Silvanus <jannik.silvanus at amd.com>
Date: Thu, 4 Jan 2024 16:58:06 +0100
Subject: [PATCH 1/2] [IR] Allow type change in ValueAsMetadata::handleRAUW

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 checking assert only, and 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.
---
 llvm/lib/IR/Metadata.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
index 515893d079b8cb..bdfbd8829186d9 100644
--- a/llvm/lib/IR/Metadata.cpp
+++ b/llvm/lib/IR/Metadata.cpp
@@ -503,7 +503,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;

>From 0ae9688f4745e127570c9e91e49c3e5053ba5507 Mon Sep 17 00:00:00 2001
From: Jannik Silvanus <jannik.silvanus at amd.com>
Date: Wed, 10 Jan 2024 17:24:56 +0100
Subject: [PATCH 2/2] Add unit test exercising handleRAUW with type changes

---
 llvm/unittests/IR/MetadataTest.cpp | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

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 =