Fix: Distinguish CFI Metadata Checks in MergeFunctions Pass (llvm#65963)

This diff fixes an issue in the MergeFunctions pass where two different
Control Flow Integrity (CFI) metadata checks were incorrectly considered
identical. These merges would lead to runtime violations down the line
as two separate objects contained a single destructor which itself
contained checks for only one of the objects.

Here I update the comparison logic to take into account the metadata at
llvm.type.test checks. Now, only truly identical checks will be
considered for merging, thus preserving the integrity of each check.

Previous discussion: https://reviews.llvm.org/D154119

Author: oskarwirga

llvm/lib/Transforms/IPO/MergeFunctions.cpp

@@ -375,9 +375,32 @@ bool MergeFunctions::doFunctionalCheck(std::vector<WeakTrackingVH> &Worklist) {
 }
 #endif
 
+/// Check whether \p F has an intrinsic which references
+/// distinct metadata as an operand. The most common
+/// instance of this would be CFI checks for function-local types.
+static bool hasDistinctMetadataIntrinsic(const Function &F) {
+  for (const BasicBlock &BB : F) {
+    for (const Instruction &I : BB.instructionsWithoutDebug()) {
+      if (!isa<IntrinsicInst>(&I))
+        continue;
+
+      for (Value *Op : I.operands()) {
+        auto *MDL = dyn_cast<MetadataAsValue>(Op);
+        if (!MDL)
+          continue;
+        if (MDNode *N = dyn_cast<MDNode>(MDL->getMetadata()))
+          if (N->isDistinct())
+            return true;
+      }
+    }
+  }
+  return false;
+}
+
 /// Check whether \p F is eligible for function merging.
 static bool isEligibleForMerging(Function &F) {
-  return !F.isDeclaration() && !F.hasAvailableExternallyLinkage();
+  return !F.isDeclaration() && !F.hasAvailableExternallyLinkage() &&
+         !hasDistinctMetadataIntrinsic(F);
 }
 
 bool MergeFunctions::runOnModule(Module &M) {

llvm/test/Transforms/MergeFunc/cfi-function-merging.ll

@@ -0,0 +1,53 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
+;; Check the cases involving internal CFI instrumented functions where we do not expect functions to be merged.
+; RUN: opt -S -passes=mergefunc < %s | FileCheck %s
+
+; Function Attrs: nocallback nofree nosync nounwind readnone speculatable willreturn
+declare i1 @llvm.type.test(ptr, metadata) #6
+
+define internal void @A__on_zero_sharedEv(ptr noundef nonnull align 8 dereferenceable(32) %this) {
+; CHECK-LABEL: define internal void @A__on_zero_sharedEv
+; CHECK-SAME: (ptr noundef nonnull align 8 dereferenceable(32) [[THIS:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[THIS_ADDR:%.*]] = alloca ptr, align 8
+; CHECK-NEXT:    store ptr [[THIS]], ptr [[THIS_ADDR]], align 8
+; CHECK-NEXT:    [[THIS1:%.*]] = load ptr, ptr [[THIS_ADDR]], align 8
+; CHECK-NEXT:    [[VTABLE:%.*]] = load ptr, ptr [[THIS1]], align 8
+; CHECK-NEXT:    [[TMP0:%.*]] = call i1 @llvm.type.test(ptr [[VTABLE]], metadata [[META0:![0-9]+]]), !nosanitize !1
+; CHECK-NEXT:    ret void
+;
+entry:
+  %this.addr = alloca ptr, align 8
+  store ptr %this, ptr %this.addr, align 8
+  %this1 = load ptr, ptr %this.addr, align 8
+  %vtable = load ptr, ptr %this1, align 8
+  %0 = call i1 @llvm.type.test(ptr %vtable, metadata !11), !nosanitize !47
+  ret void
+}
+
+; Function Attrs: mustprogress noinline nounwind optnone uwtable
+define internal void @B__on_zero_sharedEv(ptr noundef nonnull align 8 dereferenceable(32) %this) {
+; CHECK-LABEL: define internal void @B__on_zero_sharedEv
+; CHECK-SAME: (ptr noundef nonnull align 8 dereferenceable(32) [[THIS:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[THIS_ADDR:%.*]] = alloca ptr, align 8
+; CHECK-NEXT:    store ptr [[THIS]], ptr [[THIS_ADDR]], align 8
+; CHECK-NEXT:    [[THIS1:%.*]] = load ptr, ptr [[THIS_ADDR]], align 8
+; CHECK-NEXT:    [[VTABLE:%.*]] = load ptr, ptr [[THIS1]], align 8
+; CHECK-NEXT:    [[TMP0:%.*]] = call i1 @llvm.type.test(ptr [[VTABLE]], metadata [[META2:![0-9]+]]), !nosanitize !1
+; CHECK-NEXT:    ret void
+;
+entry:
+  %this.addr = alloca ptr, align 8
+  store ptr %this, ptr %this.addr, align 8
+  %this1 = load ptr, ptr %this.addr, align 8
+  %vtable = load ptr, ptr %this1, align 8
+  %0 = call i1 @llvm.type.test(ptr %vtable, metadata !22), !nosanitize !47
+  ret void
+}
+
+!10 = !{i64 16, !11}
+!11 = distinct !{}
+!21 = !{i64 16, !22}
+!22 = distinct !{}
+!47 = !{}

llvm/test/Transforms/MergeFunc/merge-fp-intrinsics.ll

@@ -0,0 +1,28 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
+;; Make sure internal constrained FP intrinsics still merge properly
+; RUN: opt -passes=mergefunc -S < %s | FileCheck %s
+
+declare float @llvm.experimental.constrained.fadd.f32(float, float, metadata, metadata)
+
+define float @func1(float %a, float %b) {
+; CHECK-LABEL: define float @func1
+; CHECK-SAME: (float [[A:%.*]], float [[B:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = call float @llvm.experimental.constrained.fadd.f32(float [[A]], float [[B]], metadata !"round.dynamic", metadata !"fpexcept.strict")
+; CHECK-NEXT:    [[RESULT_2:%.*]] = call float @llvm.experimental.constrained.fadd.f32(float [[A]], float [[B]], metadata !"round.dynamic", metadata !"fpexcept.strict")
+; CHECK-NEXT:    ret float [[RESULT]]
+;
+  %result = call float @llvm.experimental.constrained.fadd.f32(float %a, float %b, metadata !"round.dynamic", metadata !"fpexcept.strict")
+  %result_2 = call float @llvm.experimental.constrained.fadd.f32(float %a, float %b, metadata !"round.dynamic", metadata !"fpexcept.strict")
+  ret float %result
+}
+
+define float @func2(float %a, float %b) {
+; CHECK-LABEL: define float @func2
+; CHECK-SAME: (float [[TMP0:%.*]], float [[TMP1:%.*]]) {
+; CHECK-NEXT:    [[TMP3:%.*]] = tail call float @func1(float [[TMP0]], float [[TMP1]])
+; CHECK-NEXT:    ret float [[TMP3]]
+;
+  %result = call float @llvm.experimental.constrained.fadd.f32(float %a, float %b, metadata !"round.dynamic", metadata !"fpexcept.strict")
+  %result_2 = call float @llvm.experimental.constrained.fadd.f32(float %a, float %b, metadata !"round.dynamic", metadata !"fpexcept.strict")
+  ret float %result
+}