[MergeFunctions] Intersect poison-generating flags

[nikic]

Instead of comparing the flags in FunctionComparator, intersect them in MergeFunctions. This both a) fixes miscompiles where we failed to check the poison flags for GEP and FMF and b) allows more function merging, in cases where the functions only differ in flags.

This can be extended to combine instruction-level metadata as well in the future.

Fixes #219664.

[llvmorg-github-actions]

@llvm/pr-subscribers-llvm-transforms

Author: Nikita Popov (nikic)

Changes

Instead of comparing the flags in FunctionComparator, intersect them in MergeFunctions. This both a) fixes miscompiles where we failed to check the poison flags for GEP and FMF and b) allows more function merging, in cases where the functions only differ in flags.

This can be extended to combine instruction-level metadata as well in the future.

Fixes #219664.

---

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

3 Files Affected:

• (modified) llvm/lib/Transforms/IPO/MergeFunctions.cpp (+22-22)
• (modified) llvm/lib/Transforms/Utils/FunctionComparator.cpp (-6)
• (added) llvm/test/Transforms/MergeFunc/flags.ll (+86)

diff --git a/llvm/lib/Transforms/IPO/MergeFunctions.cpp b/llvm/lib/Transforms/IPO/MergeFunctions.cpp
index 02bcceaf60ebb..5a42ecebbdc7f 100644
--- a/llvm/lib/Transforms/IPO/MergeFunctions.cpp
+++ b/llvm/lib/Transforms/IPO/MergeFunctions.cpp
@@ -271,7 +271,11 @@ class MergeFunctions {
   /// again.
   void mergeTwoFunctions(Function *F, Function *G);
 
-  void mergeInstrProfMetadataInto(Function *Dst, Function *Src);
+  /// Merge \p Src's instruction-level annotations into the corresponding
+  /// instructions of \p Dst. \p Dst is the surviving function; \p Src will be
+  /// erased or rewritten after this call.
+  /// Both functions must be structurally identical.
+  void mergeInstrAnnotations(Function *Dst, Function *Src);
 
   /// Fill PDIUnrelatedWL with instructions from the entry block that are
   /// unrelated to parameter related debug info.
@@ -303,9 +307,10 @@ class MergeFunctions {
   // If needed, replace G with an alias to F if possible, or a thunk to F if
   // profitable. Returns false if neither is the case. If \p G is not needed
   // (i.e. it is discardable and not used), \p G is removed directly.
-  // \p MergeProfile must be true when G's profile should be preserved, it is
-  // merged into F before G is erased or rewritten.
-  bool writeThunkOrAliasIfNeeded(Function *F, Function *G, bool MergeProfile);
+  // If \p MergeAnnotations is true, annotations on G such as profiling
+  // information and poison-generating flags are merged into F before G is
+  // erased or rewritten.
+  bool writeThunkOrAliasIfNeeded(Function *F, Function *G, bool MergeAnnotations);
 
   /// Replace function F with function G in the function tree.
   void replaceFunctionInTree(const FunctionNode &FN, Function *G);
@@ -908,12 +913,8 @@ static void mergeEntryCountsAndImportsInto(Function &F, Function &G) {
   F.setEntryCount(Sum, AllImports.empty() ? nullptr : &AllImports);
 }
 
-// If needed, replace G with an alias to F if possible, or a thunk to F if
-// profitable. Returns false if neither is the case. If \p G is not needed (i.e.
-// it is discardable and unused), \p G is removed directly. If \p MergeProfile
-// is set, G's profile metadata is merged into F.
 bool MergeFunctions::writeThunkOrAliasIfNeeded(Function *F, Function *G,
-                                               bool MergeProfile) {
+                                               bool MergeAnnotations) {
   bool ShouldErase =
       G->isDiscardableIfUnused() && G->use_empty() && !MergeFunctionsPDI;
   bool ShouldAlias = canCreateAliasFor(G);
@@ -922,8 +923,8 @@ bool MergeFunctions::writeThunkOrAliasIfNeeded(Function *F, Function *G,
   if (!ShouldErase && !ShouldAlias && !ShouldThunk)
     return false;
 
-  if (MergeProfile) {
-    mergeInstrProfMetadataInto(F, G);
+  if (MergeAnnotations) {
+    mergeInstrAnnotations(F, G);
     mergeEntryCountsAndImportsInto(*F, *G);
   }
 
@@ -1078,11 +1079,7 @@ static void mergeValueProfileOnInstructions(Instruction *DstI,
                     VDs.size());
 }
 
-/// Merge \p Src's instruction-level branch weights and value profile
-/// metadata into the corresponding instructions of \p Dst. \p Dst is the
-/// surviving function; \p Src will be erased or rewritten after this call.
-/// Both functions must be structurally identical.
-void MergeFunctions::mergeInstrProfMetadataInto(Function *Dst, Function *Src) {
+void MergeFunctions::mergeInstrAnnotations(Function *Dst, Function *Src) {
   const BlockFrequencyInfo &DstBFI =
       FAM.getResult<BlockFrequencyAnalysis>(*Dst);
   const BlockFrequencyInfo &SrcBFI =
@@ -1095,6 +1092,9 @@ void MergeFunctions::mergeInstrProfMetadataInto(Function *Dst, Function *Src) {
   ReversePostOrderTraversal<Function *> SrcRPOT(Src);
   for (auto [DstBB, SrcBB] : llvm::zip_equal(DstRPOT, SrcRPOT)) {
     for (auto [DstI, SrcI] : llvm::zip_equal(*DstBB, *SrcBB)) {
+      // Merge poison-generating flags.
+      DstI.andIRFlags(&SrcI);
+
       MDNode *DstProf = DstI.getMetadata(LLVMContext::MD_prof);
       MDNode *SrcProf = SrcI.getMetadata(LLVMContext::MD_prof);
       if ((DstProf && isValueProfileMD(DstProf)) ||
@@ -1164,12 +1164,12 @@ void MergeFunctions::mergeTwoFunctions(Function *F, Function *G) {
     const MaybeAlign GAlign = G->getAlign();
 
     // Merge !prof, while G still has its body.
-    writeThunkOrAliasIfNeeded(F, G, /*MergeProfile*/ true);
+    writeThunkOrAliasIfNeeded(F, G, /*MergeAnnotations=*/true);
     if (FEntryCount)
       NewF->setEntryCount(*FEntryCount);
-    // NewF becomes thunk/alias to the shared body F, it has no profile to be
-    // merged.
-    writeThunkOrAliasIfNeeded(F, NewF, /*MergeProfile*/ false);
+    // NewF becomes thunk/alias to the shared body F, it has no annotations to
+    // be merged.
+    writeThunkOrAliasIfNeeded(F, NewF, /*MergeAnnotations=*/false);
 
     if (NewFAlign || GAlign)
       F->setAlignment(std::max(NewFAlign.valueOrOne(), GAlign.valueOrOne()));
@@ -1203,14 +1203,14 @@ void MergeFunctions::mergeTwoFunctions(Function *F, Function *G) {
     // stop here and delete G. There's no need for a thunk. (See note on
     // MergeFunctionsPDI above).
     if (G->isDiscardableIfUnused() && G->use_empty() && !MergeFunctionsPDI) {
-      mergeInstrProfMetadataInto(F, G);
+      mergeInstrAnnotations(F, G);
       mergeEntryCountsAndImportsInto(*F, *G);
       G->eraseFromParent();
       ++NumFunctionsMerged;
       return;
     }
 
-    if (writeThunkOrAliasIfNeeded(F, G, /*MergeProfile*/ true))
+    if (writeThunkOrAliasIfNeeded(F, G, /*MergeAnnotations=*/true))
       ++NumFunctionsMerged;
   }
 }
diff --git a/llvm/lib/Transforms/Utils/FunctionComparator.cpp b/llvm/lib/Transforms/Utils/FunctionComparator.cpp
index 20f0f6235a870..c157fbb336aff 100644
--- a/llvm/lib/Transforms/Utils/FunctionComparator.cpp
+++ b/llvm/lib/Transforms/Utils/FunctionComparator.cpp
@@ -658,8 +658,6 @@ int FunctionComparator::cmpOperations(const Instruction *L,
 
   // Differences from Instruction::isSameOperationAs:
   //  * replace type comparison with calls to cmpTypes.
-  //  * we test for I->getRawSubclassOptionalData (nuw/nsw/tail) at the top.
-  //  * because of the above, we don't test for the tail bit on calls later on.
   if (int Res = cmpNumbers(L->getOpcode(), R->getOpcode()))
     return Res;
 
@@ -678,10 +676,6 @@ int FunctionComparator::cmpOperations(const Instruction *L,
   if (int Res = cmpTypes(L->getType(), R->getType()))
     return Res;
 
-  if (int Res = cmpNumbers(L->getRawSubclassOptionalData(),
-                           R->getRawSubclassOptionalData()))
-    return Res;
-
   // We have two instructions of identical opcode and #operands.  Check to see
   // if all operands are the same type
   for (unsigned i = 0, e = L->getNumOperands(); i != e; ++i) {
diff --git a/llvm/test/Transforms/MergeFunc/flags.ll b/llvm/test/Transforms/MergeFunc/flags.ll
new file mode 100644
index 0000000000000..17907ab1f227a
--- /dev/null
+++ b/llvm/test/Transforms/MergeFunc/flags.ll
@@ -0,0 +1,86 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -passes=mergefunc < %s | FileCheck %s
+
+define i32 @fn_add_nuw_nsw(i32 %a) {
+; CHECK-LABEL: define i32 @fn_add_nuw_nsw(
+; CHECK-SAME: i32 [[A:%.*]]) {
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw i32 [[A]], 1
+; CHECK-NEXT:    ret i32 [[ADD]]
+;
+  %add = add nuw nsw i32 %a, 1
+  ret i32 %add
+}
+
+define internal i32 @fn_add_nuw(i32 %a) {
+  %add = add nuw i32 %a, 1
+  ret i32 %add
+}
+
+define ptr @fn_gep_inbounds(ptr %a) {
+; CHECK-LABEL: define ptr @fn_gep_inbounds(
+; CHECK-SAME: ptr [[A:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr nusw i8, ptr [[A]], i64 1
+; CHECK-NEXT:    ret ptr [[GEP]]
+;
+  %gep = getelementptr inbounds i8, ptr %a, i64 1
+  ret ptr %gep
+}
+
+define internal ptr @fn_gep_nusw(ptr %a) {
+  %gep = getelementptr nusw i8, ptr %a, i64 1
+  ret ptr %gep
+}
+
+define ptr @fn_gep_inbounds2(ptr %a) {
+; CHECK-LABEL: define ptr @fn_gep_inbounds2(
+; CHECK-SAME: ptr [[A:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[A]], i64 2
+; CHECK-NEXT:    ret ptr [[GEP]]
+;
+  %gep = getelementptr inbounds i8, ptr %a, i64 2
+  ret ptr %gep
+}
+
+define internal ptr @fn_gep_nuw(ptr %a) {
+  %gep = getelementptr nuw i8, ptr %a, i64 2
+  ret ptr %gep
+}
+
+define float @fn_fadd_ninf_nnan(float %a) {
+; CHECK-LABEL: define float @fn_fadd_ninf_nnan(
+; CHECK-SAME: float [[A:%.*]]) {
+; CHECK-NEXT:    [[ADD:%.*]] = fadd ninf float [[A]], 1.000000e+00
+; CHECK-NEXT:    ret float [[ADD]]
+;
+  %add = fadd ninf nnan float %a, 1.0
+  ret float %add
+}
+
+define internal float @fn_fadd_ninf_reassoc(float %a) {
+  %add = fadd ninf reassoc float %a, 1.0
+  ret float %add
+}
+
+define void @calls(i32 %x, ptr %p, float %f) {
+; CHECK-LABEL: define void @calls(
+; CHECK-SAME: i32 [[X:%.*]], ptr [[P:%.*]], float [[F:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @fn_add_nuw_nsw(i32 [[X]])
+; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @fn_add_nuw_nsw(i32 [[X]])
+; CHECK-NEXT:    [[TMP3:%.*]] = call ptr @fn_gep_inbounds(ptr [[P]])
+; CHECK-NEXT:    [[TMP4:%.*]] = call ptr @fn_gep_inbounds(ptr [[P]])
+; CHECK-NEXT:    [[TMP5:%.*]] = call ptr @fn_gep_inbounds2(ptr [[P]])
+; CHECK-NEXT:    [[TMP6:%.*]] = call ptr @fn_gep_inbounds2(ptr [[P]])
+; CHECK-NEXT:    [[TMP7:%.*]] = call float @fn_fadd_ninf_nnan(float [[F]])
+; CHECK-NEXT:    [[TMP8:%.*]] = call float @fn_fadd_ninf_nnan(float [[F]])
+; CHECK-NEXT:    ret void
+;
+  call i32 @fn_add_nuw_nsw(i32 %x)
+  call i32 @fn_add_nuw(i32 %x)
+  call ptr @fn_gep_inbounds(ptr %p)
+  call ptr @fn_gep_nusw(ptr %p)
+  call ptr @fn_gep_inbounds2(ptr %p)
+  call ptr @fn_gep_nuw(ptr %p)
+  call float @fn_fadd_ninf_nnan(float %f)
+  call float @fn_fadd_ninf_reassoc(float %f)
+  ret void
+}

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[antoniofrighetto]

LGTM

[antoniofrighetto]

Minor nit: the comment was quite stale, though perhaps we could still mention that we do test the tail bit on calls later (CallBase branch)?

[nikic]

I went ahead and just dropped this comment entirely. I don't think it really provides useful information.