[PseudoProbe] Clean up dwarf discriminator and avoid duplicating factor.

A pseudo probe is created with dwarf line information shared with its nearest instruction. If the instruction comes with a dwarf discriminator, it will be shared with the probe as well. This can confuse the later FS-AFDO discriminator assignment pass. To fix this, I'm cleaning up the discriminator fields for probes when they are inserted.

I also notice another possibility to change the discriminator field of pseudo probes in the pipeline before the FS discriminator assignment pass. That is the loop unroller, which assigns duplication factor to instruction being vectorized. I'm disabling that for pseudo probe intrinsics specifically, also for callsites with probes.

Reviewed By: wenlei

Differential Revision: https://reviews.llvm.org/D148569

llvm/include/llvm/IR/DebugInfoMetadata.h

@@ -2308,6 +2308,12 @@ DILocation::cloneWithBaseDiscriminator(unsigned D) const {
 std::optional<const DILocation *>
 DILocation::cloneByMultiplyingDuplicationFactor(unsigned DF) const {
   assert(!EnableFSDiscriminator && "FSDiscriminator should not call this.");
+  // Do no interfere with pseudo probes. Pseudo probe doesn't need duplication
+  // factor support as samples collected on cloned probes will be aggregated.
+  // Also pseudo probe at a callsite uses the dwarf discriminator to store
+  // pseudo probe related information, such as the probe id.
+  if (isPseudoProbeDiscriminator(getDiscriminator()))
+    return this;
 
   DF *= getDuplicationFactor();
   if (DF <= 1)

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

@@ -349,6 +349,14 @@ void SampleProfileProber::instrumentOneFunc(Function &F, TargetMachine *TM) {
                      Builder.getInt64(PseudoProbeFullDistributionFactor)};
     auto *Probe = Builder.CreateCall(ProbeFn, Args);
     AssignDebugLoc(Probe);
+    // Reset the dwarf discriminator if the debug location comes with any. The
+    // discriminator field may be used by FS-AFDO later in the pipeline.
+    if (auto DIL = Probe->getDebugLoc()) {
+      if (DIL->getDiscriminator()) {
+        DIL = DIL->cloneWithDiscriminator(0);
+        Probe->setDebugLoc(DIL);
+      }
+    }
   }
 
   // Probe both direct calls and indirect calls. Direct calls are probed so that
@@ -361,12 +369,13 @@ void SampleProfileProber::instrumentOneFunc(Function &F, TargetMachine *TM) {
                         ? (uint32_t)PseudoProbeType::DirectCall
                         : (uint32_t)PseudoProbeType::IndirectCall;
     AssignDebugLoc(Call);
-    // Levarge the 32-bit discriminator field of debug data to store the ID and
-    // type of a callsite probe. This gets rid of the dependency on plumbing a
-    // customized metadata through the codegen pipeline.
-    uint32_t V = PseudoProbeDwarfDiscriminator::packProbeData(
-        Index, Type, 0, PseudoProbeDwarfDiscriminator::FullDistributionFactor);
     if (auto DIL = Call->getDebugLoc()) {
+      // Levarge the 32-bit discriminator field of debug data to store the ID
+      // and type of a callsite probe. This gets rid of the dependency on
+      // plumbing a customized metadata through the codegen pipeline.
+      uint32_t V = PseudoProbeDwarfDiscriminator::packProbeData(
+          Index, Type, 0,
+          PseudoProbeDwarfDiscriminator::FullDistributionFactor);
       DIL = DIL->cloneWithDiscriminator(V);
       Call->setDebugLoc(DIL);
     }

llvm/lib/Transforms/Utils/LoopUnroll.cpp

@@ -516,7 +516,7 @@ LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
       !EnableFSDiscriminator)
     for (BasicBlock *BB : L->getBlocks())
       for (Instruction &I : *BB)
-        if (!isa<DbgInfoIntrinsic>(&I))
+        if (!I.isDebugOrPseudoInst())
           if (const DILocation *DIL = I.getDebugLoc()) {
             auto NewDIL = DIL->cloneByMultiplyingDuplicationFactor(ULO.Count);
             if (NewDIL)

llvm/lib/Transforms/Utils/LoopUnrollAndJam.cpp

@@ -347,7 +347,7 @@ llvm::UnrollAndJamLoop(Loop *L, unsigned Count, unsigned TripCount,
       !EnableFSDiscriminator)
     for (BasicBlock *BB : L->getBlocks())
       for (Instruction &I : *BB)
-        if (!isa<DbgInfoIntrinsic>(&I))
+        if (!I.isDebugOrPseudoInst())
           if (const DILocation *DIL = I.getDebugLoc()) {
             auto NewDIL = DIL->cloneByMultiplyingDuplicationFactor(Count);
             if (NewDIL)

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -277,7 +277,7 @@ void VPTransformState::setDebugLocFromInst(const Value *V) {
   // When a FSDiscriminator is enabled, we don't need to add the multiply
   // factors to the discriminators.
   if (DIL && Inst->getFunction()->shouldEmitDebugInfoForProfiling() &&
-      !isa<DbgInfoIntrinsic>(Inst) && !EnableFSDiscriminator) {
+      !Inst->isDebugOrPseudoInst() && !EnableFSDiscriminator) {
     // FIXME: For scalable vectors, assume vscale=1.
     auto NewDIL =
         DIL->cloneByMultiplyingDuplicationFactor(UF * VF.getKnownMinValue());

llvm/test/Transforms/LoopVectorize/discriminator.ll

@@ -2,6 +2,7 @@
 ; RUN: opt -S -passes=loop-vectorize -force-vector-width=2 -force-vector-interleave=3 < %s | FileCheck --check-prefix=DBG_VALUE --check-prefix=LOOPVEC_2_3 %s
 ; RUN: opt -S -passes=loop-unroll  -unroll-count=5 < %s | FileCheck --check-prefix=DBG_VALUE --check-prefix=LOOPUNROLL_5 %s
 ; RUN: opt -S -passes=loop-vectorize,loop-unroll -force-vector-width=4 -force-vector-interleave=4 - -unroll-count=2 < %s | FileCheck --check-prefix=DBG_VALUE --check-prefix=LOOPVEC_UNROLL %s
+; RUN: opt -S -passes=pseudo-probe,loop-unroll  -unroll-count=5 < %s | FileCheck --check-prefix=PSEUDO_PROBE %s
 
 ; Test if vectorization/unroll factor is recorded in discriminator.
 ;
@@ -30,6 +31,7 @@ define void @_Z3foov() local_unnamed_addr #0 !dbg !6 {
   %6 = getelementptr inbounds i32, ptr %2, i64 %indvars.iv, !dbg !13
   %7 = load i32, ptr %6, align 4, !dbg !17, !tbaa !15
   %8 = add nsw i32 %7, %5, !dbg !17
+;PSEUDO_PROBE-COUNT-5: call void @llvm.pseudoprobe(i64 6699318081062747564, i64 2, i32 0, i64 -1), !dbg ![[#PROBE:]]
 ;DBG_VALUE: call void @llvm.dbg.declare{{.*}}!dbg ![[DBG:[0-9]*]]
   call void @llvm.dbg.declare(metadata i32 %8, metadata !22, metadata !DIExpression()), !dbg !17
   store i32 %8, ptr %6, align 4, !dbg !17, !tbaa !15
@@ -50,6 +52,9 @@ define void @_Z3foov() local_unnamed_addr #0 !dbg !6 {
 ;LOOPVEC_UNROLL: discriminator: 9
 ;LOOPVEC_UNROLL: discriminator: 385
 ;DBG_VALUE: ![[DBG]] = {{.*}}, scope: ![[TOP]]
+; Pseudo probe should not have duplication factor assigned.
+;PSEUDO_PROBE: ![[TOP:[0-9]*]] = distinct !DISubprogram(name: "foo"
+;PSEUDO_PROBE: ![[#PROBE]] = !DILocation(line: 6, column: 13, scope: ![[TOP]])
 
 !llvm.dbg.cu = !{!0}
 !llvm.module.flags = !{!3, !4}

llvm/test/Transforms/SampleProfile/pseudo-probe-emit.ll

@@ -69,7 +69,8 @@ entry:
 
 ; CHECK-IL: ![[#FOO:]] = distinct !DISubprogram(name: "foo"
 ; CHECK-IL: ![[#FAKELINE]] = !DILocation(line: 0, scope: ![[#FOO]])
-; CHECK-IL: ![[#REALLINE]] = !DILocation(line: 2, scope: ![[#FOO]])
+; CHECK-IL: ![[#REALLINE]] = !DILocation(line: 2, scope: ![[#DISC0:]])
+; CHECK-IL: ![[#DISC0]] = !DILexicalBlockFile(scope: ![[#FOO]], file: ![[#]], discriminator: 0)
 ; CHECK-IL: ![[#PROBE0]] = !DILocation(line: 2, column: 20, scope: ![[#SCOPE0:]])
 ;; A discriminator of 67108887 which is 0x7200017 in hexdecimal, stands for a direct call probe
 ;; with an index of 2.
@@ -109,5 +110,6 @@ entry:
 !9 = !{i32 2, !"Dwarf Version", i32 4}
 !10 = !{i32 2, !"Debug Info Version", i32 3}
 !11 = !{!"clang version 3.9.0"}
-!12 = !DILocation(line: 2, scope: !3)
+!12 = !DILocation(line: 2, scope: !14)
 !13 = !DILocation(line: 2, column: 20, scope: !4)
+!14 = !DILexicalBlockFile(scope: !3, file: !1, discriminator: 1)

llvm/test/Transforms/SampleProfile/pseudo-probe-profile-mismatch.ll

@@ -232,24 +232,24 @@ attributes #5 = { nocallback nofree nosync nounwind readnone speculatable willre
 !60 = !DILocation(line: 0, scope: !54)
 !61 = !DILocation(line: 16, column: 8, scope: !54)
 !62 = !DILocation(line: 16, column: 19, scope: !63)
-!63 = !DILexicalBlockFile(scope: !58, file: !3, discriminator: 2)
+!63 = !DILexicalBlockFile(scope: !58, file: !3, discriminator: 0)
 !64 = !DILocation(line: 16, column: 21, scope: !63)
 !65 = !DILocation(line: 16, column: 3, scope: !66)
-!66 = !DILexicalBlockFile(scope: !54, file: !3, discriminator: 2)
+!66 = !DILexicalBlockFile(scope: !54, file: !3, discriminator: 0)
 !67 = !DILocation(line: 0, scope: !49)
 !68 = !DILocation(line: 23, column: 1, scope: !49)
 !69 = !DILocation(line: 17, column: 14, scope: !56)
 !70 = !DILocation(line: 0, scope: !56)
 !71 = !DILocation(line: 17, column: 10, scope: !56)
 !72 = !DILocation(line: 17, column: 21, scope: !73)
-!73 = !DILexicalBlockFile(scope: !74, file: !3, discriminator: 2)
+!73 = !DILexicalBlockFile(scope: !74, file: !3, discriminator: 0)
 !74 = distinct !DILexicalBlock(scope: !56, file: !3, line: 17, column: 5)
 !75 = !DILocation(line: 17, column: 23, scope: !73)
 !76 = !DILocation(line: 17, column: 5, scope: !77)
-!77 = !DILexicalBlockFile(scope: !56, file: !3, discriminator: 2)
+!77 = !DILexicalBlockFile(scope: !56, file: !3, discriminator: 0)
 !78 = !DILocation(line: 22, column: 3, scope: !57)
 !79 = !DILocation(line: 16, column: 30, scope: !80)
-!80 = !DILexicalBlockFile(scope: !58, file: !3, discriminator: 4)
+!80 = !DILexicalBlockFile(scope: !58, file: !3, discriminator: 0)
 !81 = !DILocation(line: 16, column: 3, scope: !80)
 !82 = distinct !{!82, !83, !84, !85}
 !83 = !DILocation(line: 16, column: 3, scope: !54)
@@ -269,7 +269,7 @@ attributes #5 = { nocallback nofree nosync nounwind readnone speculatable willre
 !97 = !DILexicalBlockFile(scope: !87, file: !3, discriminator: 186646647)
 !98 = !DILocation(line: 20, column: 9, scope: !87)
 !99 = !DILocation(line: 17, column: 33, scope: !100)
-!100 = !DILexicalBlockFile(scope: !74, file: !3, discriminator: 4)
+!100 = !DILexicalBlockFile(scope: !74, file: !3, discriminator: 0)
 !101 = !DILocation(line: 17, column: 5, scope: !100)
 !102 = distinct !{!102, !103, !104, !85}
 !103 = !DILocation(line: 17, column: 5, scope: !56)