[Metadata] Preserve MD_prof when merging instructions when one is missing.

[snehasish]

Preserve branch weight metadata when merging instructions if one of the
instructions is missing metadata. This is similar in behaviour to what
we do today for other types of metadata such as mmra, memprof and
callsite metadata.

[snehasish]

• [Metadata] Preserve MD_prof when merging instructions when one is missing. #132433 👈 (View in Graphite)
• [Metadata] Handle memprof, callsite merging when one is missing. #132106
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Snehasish Kumar (snehasish)

Changes

Preserve branch weight metadata when merging instructions if one of the
instructions is missing metadata. This is similar in behaviour to what
we do today for other types of metadata such as mmra, memprof and
callsite metadata.

---

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

3 Files Affected:

• (modified) llvm/lib/Transforms/Utils/Local.cpp (+13-5)
• (added) llvm/test/Transforms/SimplifyCFG/merge-direct-call-branch-weights-preserve-hoist.ll (+62)
• (added) llvm/test/Transforms/SimplifyCFG/merge-direct-call-branch-weights-preserve-sink.ll (+62)

diff --git a/llvm/lib/Transforms/Utils/Local.cpp b/llvm/lib/Transforms/Utils/Local.cpp
index d18adac5fa914..069ebae117b48 100644
--- a/llvm/lib/Transforms/Utils/Local.cpp
+++ b/llvm/lib/Transforms/Utils/Local.cpp
@@ -3355,9 +3355,11 @@ static void combineMetadata(Instruction *K, const Instruction *J,
       case LLVMContext::MD_invariant_group:
         // Preserve !invariant.group in K.
         break;
-      // Keep empty cases for mmra, memprof, and callsite to prevent them from
+      // Keep empty cases for prof, mmra, memprof, and callsite to prevent them from
       // being removed as unknown metadata. The actual merging is handled
       // separately below.
+      case LLVMContext::MD_prof:
+        [[fallthrough]];
       case LLVMContext::MD_mmra:
       case LLVMContext::MD_memprof:
       case LLVMContext::MD_callsite:
@@ -3386,10 +3388,6 @@ static void combineMetadata(Instruction *K, const Instruction *J,
         if (!AAOnly)
           K->setMetadata(Kind, JMD);
         break;
-      case LLVMContext::MD_prof:
-        if (!AAOnly && DoesKMove)
-          K->setMetadata(Kind, MDNode::getMergedProfMetadata(KMD, JMD, K, J));
-        break;
       case LLVMContext::MD_noalias_addrspace:
         if (DoesKMove)
           K->setMetadata(Kind,
@@ -3436,6 +3434,16 @@ static void combineMetadata(Instruction *K, const Instruction *J,
     K->setMetadata(LLVMContext::MD_callsite,
                    MDNode::getMergedCallsiteMetadata(KCallSite, JCallSite));
   }
+
+  // Merge prof metadata.
+  // Handle separately to support cases where only one instruction has the
+  // metadata.
+  auto JProf = J->getMetadata(LLVMContext::MD_prof);
+  auto KProf = K->getMetadata(LLVMContext::MD_prof);
+  if (!AAOnly && (JProf || KProf)) {
+    K->setMetadata(LLVMContext::MD_prof,
+                   MDNode::getMergedProfMetadata(KProf, JProf, K, J));
+  }
 }
 
 void llvm::combineMetadataForCSE(Instruction *K, const Instruction *J,
diff --git a/llvm/test/Transforms/SimplifyCFG/merge-direct-call-branch-weights-preserve-hoist.ll b/llvm/test/Transforms/SimplifyCFG/merge-direct-call-branch-weights-preserve-hoist.ll
new file mode 100644
index 0000000000000..2a0fc3ffc2d4d
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/merge-direct-call-branch-weights-preserve-hoist.ll
@@ -0,0 +1,62 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals --version 2
+; RUN: opt < %s -passes='simplifycfg<no-sink-common-insts;hoist-common-insts>' -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck %s --check-prefix=HOIST
+
+; Test case based on C++ code with manualy annotated !prof metadata.
+; This is to test that when calls to 'func1' from 'if.then' block
+; and 'if.else' block are hoisted, the branch_weights are merged and
+; attached to merged call rather than dropped.
+;
+; int func1(int a, int b) ;
+; int func2(int a, int b) ;
+
+; int func(int a, int b, bool c) {
+;    int sum= 0;
+;    if(c) {
+;        sum += func1(a, b);
+;    } else {
+;        sum += func1(a, b);
+;        sum -= func2(a, b);
+;    }
+;    return sum;
+; }
+define i32 @_Z4funciib(i32 %a, i32 %b, i1 %c) {
+; HOIST-LABEL: define i32 @_Z4funciib
+; HOIST-SAME: (i32 [[A:%.*]], i32 [[B:%.*]], i1 [[C:%.*]]) {
+; HOIST-NEXT:  entry:
+; HOIST-NEXT:    [[CALL:%.*]] = tail call i32 @_Z5func1ii(i32 [[A]], i32 [[B]]), !prof [[PROF0:![0-9]+]]
+; HOIST-NEXT:    br i1 [[C]], label [[IF_END:%.*]], label [[IF_ELSE:%.*]]
+; HOIST:       if.else:
+; HOIST-NEXT:    [[CALL3:%.*]] = tail call i32 @_Z5func2ii(i32 [[A]], i32 [[B]])
+; HOIST-NEXT:    [[SUB:%.*]] = sub i32 [[CALL]], [[CALL3]]
+; HOIST-NEXT:    br label [[IF_END]]
+; HOIST:       if.end:
+; HOIST-NEXT:    [[SUM_0:%.*]] = phi i32 [ [[SUB]], [[IF_ELSE]] ], [ [[CALL]], [[ENTRY:%.*]] ]
+; HOIST-NEXT:    ret i32 [[SUM_0]]
+;
+entry:
+  br i1 %c, label %if.then, label %if.else
+
+if.then:                                          ; preds = %entry
+  %call = tail call i32 @_Z5func1ii(i32 %a, i32 %b), !prof !0
+  br label %if.end
+
+if.else:                                          ; preds = %entry
+  %call1 = tail call i32 @_Z5func1ii(i32 %a, i32 %b)
+  %call3 = tail call i32 @_Z5func2ii(i32 %a, i32 %b)
+  %sub = sub i32 %call1, %call3
+  br label %if.end
+
+if.end:                                           ; preds = %if.else, %if.then
+  %sum.0 = phi i32 [ %call, %if.then ], [ %sub, %if.else ]
+  ret i32 %sum.0
+}
+
+declare i32 @_Z5func1ii(i32, i32)
+
+declare i32 @_Z5func2ii(i32, i32)
+
+!0 = !{!"branch_weights", i32 10}
+!1 = !{!"branch_weights", i32 90}
+;.
+; HOIST: [[PROF0]] = !{!"branch_weights", i32 10}
+;.
diff --git a/llvm/test/Transforms/SimplifyCFG/merge-direct-call-branch-weights-preserve-sink.ll b/llvm/test/Transforms/SimplifyCFG/merge-direct-call-branch-weights-preserve-sink.ll
new file mode 100644
index 0000000000000..6ae798abce827
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/merge-direct-call-branch-weights-preserve-sink.ll
@@ -0,0 +1,62 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals --version 2
+; RUN: opt < %s -passes='simplifycfg<sink-common-insts;no-hoist-common-insts>' -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck %s --check-prefix=SINK
+
+
+; Test case based on the following C++ code with manualy annotated !prof metadata.
+; This is to test that when calls to 'func1' from 'if.then' and 'if.else' are
+; sinked, the branch weights are merged and attached to sinked call.
+;
+; int func1(int a, int b) ;
+; int func2(int a, int b) ;
+
+; int func(int a, int b, bool c) {
+;    int sum = 0;
+;    if (c) {
+;        sum += func1(a,b);
+;    } else {
+;        b -= func2(a,b);
+;        sum += func1(a,b);
+;    }
+;    return sum;
+; }
+
+define i32 @_Z4funciib(i32 %a, i32 %b, i1 %c) {
+; SINK-LABEL: define i32 @_Z4funciib
+; SINK-SAME: (i32 [[A:%.*]], i32 [[B:%.*]], i1 [[C:%.*]]) {
+; SINK-NEXT:  entry:
+; SINK-NEXT:    br i1 [[C]], label [[IF_END:%.*]], label [[IF_ELSE:%.*]]
+; SINK:       if.else:
+; SINK-NEXT:    [[CALL1:%.*]] = tail call i32 @_Z5func2ii(i32 [[A]], i32 [[B]])
+; SINK-NEXT:    [[SUB:%.*]] = sub i32 [[B]], [[CALL1]]
+; SINK-NEXT:    br label [[IF_END]]
+; SINK:       if.end:
+; SINK-NEXT:    [[SUB_SINK:%.*]] = phi i32 [ [[SUB]], [[IF_ELSE]] ], [ [[B]], [[ENTRY:%.*]] ]
+; SINK-NEXT:    [[CALL2:%.*]] = tail call i32 @_Z5func1ii(i32 [[A]], i32 [[SUB_SINK]]), !prof [[PROF0:![0-9]+]]
+; SINK-NEXT:    ret i32 [[CALL2]]
+;
+entry:
+  br i1 %c, label %if.then, label %if.else
+
+if.then:                                          ; preds = %entry
+  %call = tail call i32 @_Z5func1ii(i32 %a, i32 %b)
+  br label %if.end
+
+if.else:                                          ; preds = %entry
+  %call1 = tail call i32 @_Z5func2ii(i32 %a, i32 %b)
+  %sub = sub i32 %b, %call1
+  %call2 = tail call i32 @_Z5func1ii(i32 %a, i32 %sub), !prof !0
+  br label %if.end
+
+if.end:                                           ; preds = %if.else, %if.then
+  %sum.0 = phi i32 [ %call, %if.then ], [ %call2, %if.else ]
+  ret i32 %sum.0
+}
+
+declare i32 @_Z5func1ii(i32, i32)
+
+declare i32 @_Z5func2ii(i32, i32)
+
+!0 = !{!"branch_weights", i32 10}
+;.
+; SINK: [[PROF0]] = !{!"branch_weights", i32 10}
+;.

[github-actions]

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

[teresajohnson]

The old handling was guarded on DoesKMove - what is the implication of removing that?

[snehasish]

Removing the condition was intentional. In the case that DoesKMove is true, the merging of md_prof can be skipped if J does not have prof metadata (a minor optimization). I felt that it was simpler to just perform the merge regardless. Let me know if you feel otherwise.

[snehasish]

Merge activity

• Apr 2, 1:06 PM EDT: A user started a stack merge that includes this pull request via Graphite.
• Apr 2, 1:11 PM EDT: Graphite rebased this pull request as part of a merge.
• Apr 2, 1:13 PM EDT: A user merged this pull request with Graphite.