[PGO][indirect-call-promotion] Add extra parameter check for musttail callsite

[xur]

Summary

Add extra parameter type check to musttail callsite to avoid verifier error.

[rnk]

llvm/test/Transforms/PGOProfile/indirect_call_promotion_musttail_typecheck.ll

30

I would like to see more of the successfully transformed IR, I want to see if the musttail call is in the tail position, since that's also a verifier requirement.

Can we use https://github.com/llvm/llvm-project/blob/main/llvm/utils/update_test_checks.py for this test? There's not too much extraneous IR.

[xur]

Followed Reid's suggestion to explicitly check the IR after the pass.

[tejohnson]

lgtm with a couple comment suggestions for test

llvm/test/Transforms/PGOProfile/indirect_call_promotion_musttail_typecheck.ll

13

Add comment that we are checking that there was no icp due to parameter mismatch.

34

Add comment that we are checking that icp succeeded since the parameter types match and that the promoted call is also musttail.

[rnk]

Thanks, my concerns are addressed!

[xur]

Add comments in the tests suggested by Teresa.

llvm/lib/Transforms/Utils/CallPromotionUtils.cpp

for (; I < NumParams; ++I) {
     Type *ActualTy = CB.getArgOperand(I)->getType();
     if (FormalTy == ActualTy)
       continue;
     if (!CastInst::isBitOrNoopPointerCastable(ActualTy, FormalTy, DL)) {
       if (FailureReason)
         *FailureReason = "Argument type mismatch";
       return false;
     }
+
+    // MustTail call needs stricter type match. See
+    // Verifier::verifyMustTailCall().
+    if (CB.isMustTailCall()) {
+      PointerType *PF = dyn_cast<PointerType>(FormalTy);
+      PointerType *PA = dyn_cast<PointerType>(ActualTy);
+      if (!PF || !PA || PF->getAddressSpace() != PA->getAddressSpace()) {
+        if (FailureReason)
+          *FailureReason = "Musttail call Argument type mismatch";
+        return false;
+      }
+    }
   }
   for (; I < NumArgs; I++) {
     // Vararg functions can have more arguments than parameters.
     assert(Callee->isVarArg());
     if (CB.paramHasAttr(I, Attribute::StructRet)) {
       if (FailureReason)
         *FailureReason = "SRet arg to vararg function";
       return false;

llvm/test/Transforms/PGOProfile/indirect_call_promotion_musttail_typecheck.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; Mustcall needs stricter parameter type checks otherwise it will fail in verifier.
+
+; RUN: opt < %s -passes=pgo-icall-prom -S | FileCheck %s
+
+; Here we check there is no ICP due to parameter mismatch.
+define ptr @func(ptr %msg, ptr %ptr, ptr %ctx, i64 %data.coerce, ptr %table, i64 %hasbits) {
+; CHECK-LABEL: @func(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = load ptr, ptr null, align 8
+; CHECK-NEXT:    ret ptr null
+; CHECK:       1:
+; CHECK-NEXT:    [[CALL11_I:%.*]] = musttail call ptr [[TMP0]](ptr null, ptr null, ptr null, i64 0, ptr null, i64 0), !prof [[PROF0:![0-9]+]]
+; CHECK-NEXT:    ret ptr [[CALL11_I]]
+;
+entry:
+  %0 = load ptr, ptr null, align 8
+  ret ptr null
+
+1:
+  %call11.i = musttail call ptr %0(ptr null, ptr null, ptr null, i64 0, ptr null, i64 0), !prof !0
+  ret ptr %call11.i
+}
+
+; Here we check that ICP succeeds since parameters match. Also check the direct-call has a MustCall attribute.
+define ptr @func2(ptr %msg, i64 %tag, ptr %ctx, ptr %type, ptr %table, ptr %ptr) {
+; CHECK-LABEL: @func2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP1:%.*]] = load ptr, ptr null, align 8
+; CHECK-NEXT:    ret ptr null
+; CHECK:       1:
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq ptr [[TMP1]], @_ZN6proto28internal12ExtensionSet10ParseFieldEmPKcPKNS_7MessageEPNS0_16InternalMetadataEPNS0_12ParseContextE
+; CHECK-NEXT:    br i1 [[TMP2]], label [[IF_TRUE_DIRECT_TARG:%.*]], label [[TMP4:%.*]], !prof [[PROF1:![0-9]+]]
+; CHECK:       if.true.direct_targ:
+; CHECK-NEXT:    [[TMP3:%.*]] = musttail call ptr @_ZN6proto28internal12ExtensionSet10ParseFieldEmPKcPKNS_7MessageEPNS0_16InternalMetadataEPNS0_12ParseContextE(ptr null, i64 0, ptr null, ptr null, ptr null, ptr null)
+; CHECK-NEXT:    ret ptr [[TMP3]]
+; CHECK:       4:
+; CHECK-NEXT:    [[CALL11_I:%.*]] = musttail call ptr [[TMP1]](ptr null, i64 0, ptr null, ptr null, ptr null, ptr null), !prof [[PROF2:![0-9]+]]
+; CHECK-NEXT:    ret ptr [[CALL11_I]]
+;
+entry:
+  %0 = load ptr, ptr null, align 8
+  ret ptr null
+
+1:
+  %call11.i = musttail call ptr %0(ptr null, i64 0, ptr null, ptr null, ptr null, ptr null), !prof !0
+  ret ptr %call11.i
+}
+
+define available_externally ptr @_ZN6proto28internal12ExtensionSet10ParseFieldEmPKcPKNS_7MessageEPNS0_16InternalMetadataEPNS0_12ParseContextE(ptr %this, i64 %tag, ptr %ptr, ptr %containing_type, ptr %metadata, ptr %ctx) {
+entry:
+  ret ptr null
+}
+
+!0 = !{!"VP", i32 0, i64 2024, i64 -4843250054591211088, i64 -1, i64 1456131869974120143, i64 947, i64 -4941069334091589447, i64 18}