[FunctionSpecialization] Fix profile count preserving logic (#157939)

The previous fix in #157768 had a bug; instead of subtracting the
original function's call count per call site of a specialization, we
were subtracting the running total of the specialization's calls.

Tracking issue: #147390

Author: alanzhao1

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

@@ -802,8 +802,8 @@ bool FunctionSpecializer::run() {
         if (std::optional<llvm::Function::ProfileCount> MaybeOriginalCount =
                 S.F->getEntryCount()) {
           uint64_t OriginalCount = MaybeOriginalCount->getCount();
-          if (OriginalCount >= CallCount) {
-            S.F->setEntryCount(OriginalCount - CallCount);
+          if (OriginalCount >= *Count) {
+            S.F->setEntryCount(OriginalCount - *Count);
           } else {
             // This should generally not happen as that would mean there are
             // more computed calls to the function than what was recorded.

llvm/test/Transforms/FunctionSpecialization/profile-counts.ll

@@ -13,23 +13,29 @@ entry:
 
 ; CHECK: if.then:
 ; CHECK: call i32 @foo.specialized.1(i32 %x, ptr @A)
+; CHECK: call i32 @foo.specialized.2(i32 %y, ptr @B)
+; CHECK: call i32 @foo.specialized.2(i32 %y, ptr @B)
 if.then:
   %call = call i32 @foo(i32 %x, ptr @A)
+  %call1 = call i32 @foo(i32 %y, ptr @B)
+  %call2 = call i32 @foo(i32 %y, ptr @B)
+  %add = add i32 %call, %call1
+  %add1 = add i32 %add, %call2
   br label %return
 
 ; CHECK: if.else:
 ; CHECK: call i32 @foo.specialized.2(i32 %y, ptr @B)
 if.else:
-  %call1 = call i32 @foo(i32 %y, ptr @B)
+  %call3 = call i32 @foo(i32 %y, ptr @B)
   br label %return
 
 ; CHECK: return:
-; CHECK: %call2 = call i32 @foo(i32 %x, ptr %z)
+; CHECK: call i32 @foo(i32 %x, ptr %z)
 return:
-  %retval.0 = phi i32 [ %call, %if.then ], [ %call1, %if.else ]
-  %call2 = call i32 @foo(i32 %x, ptr %z);
-  %add = add i32 %retval.0, %call2
-  ret i32 %add
+  %retval.0 = phi i32 [ %add1, %if.then ], [ %call3, %if.else ]
+  %call4 = call i32 @foo(i32 %x, ptr %z);
+  %add2 = add i32 %retval.0, %call4
+  ret i32 %add2
 }
 
 ; CHECK: define internal i32 @foo(i32 %x, ptr %b) !prof ![[FOO_UNSPEC_PROF:[0-9]]]
@@ -44,9 +50,9 @@ entry:
 
 ; CHECK: ![[BAR_PROF]] = !{!"function_entry_count", i64 1000}
 ; CHECK: ![[BRANCH_PROF]] = !{!"branch_weights", i32 1, i32 3}
-; CHECK: ![[FOO_UNSPEC_PROF]] =  !{!"function_entry_count", i64 234}
+; CHECK: ![[FOO_UNSPEC_PROF]] =  !{!"function_entry_count", i64 500}
 ; CHECK: ![[FOO_SPEC_1_PROF]] = !{!"function_entry_count", i64 250}
-; CHECK: ![[FOO_SPEC_2_PROF]] = !{!"function_entry_count", i64 750}
+; CHECK: ![[FOO_SPEC_2_PROF]] = !{!"function_entry_count", i64 1250}
 !0 = !{!"function_entry_count", i64 1000}
 !1 = !{!"branch_weights", i32 1, i32 3}
-!2 = !{!"function_entry_count", i64 1234}
+!2 = !{!"function_entry_count", i64 2000}

llvm/utils/profcheck-xfail.txt

@@ -755,7 +755,6 @@ Transforms/FixIrreducible/basic.ll
 Transforms/FixIrreducible/bug45623.ll
 Transforms/FixIrreducible/nested.ll
 Transforms/FixIrreducible/switch.ll
-Transforms/FunctionSpecialization/function-specialization3.ll
 Transforms/GCOVProfiling/atomic-counter.ll
 Transforms/GCOVProfiling/exit-block.ll
 Transforms/GCOVProfiling/function-numbering.ll