[CSSPGO][llvm-profgen] Add brackets for context id to support extended binary format

[wlei]

Summary

To align with https://reviews.llvm.org/D95547, we need to add brackets for context id before initializing the SampleContext.

Also added test cases for extended binary format from llvm-profgen side.

[hoy]

llvm/lib/ProfileData/SampleProfWriter.cpp

362–363

The whole if-else can be just replaced by OS << S.getNameWithContext(true) << ":" << S.getTotalSamples();

llvm/tools/llvm-profgen/ProfileGenerator.cpp

278

Should this be run unconditionally?

[wlei]

llvm/lib/ProfileData/SampleProfWriter.cpp

362–363

I see, thanks, FunctionSamples::ProfileIsCS is already checked in getNameWithContext.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

278

in line273 ContextId.rsplit(" @ ").first.str(); has already had the '[', so only add for the else branch

[wlei]

address Hongtao's feedback

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

278

I see. I'm wondering if it's possible to find a place that can centralize the appending of brackets, say right before populating or looking up the function profiles?

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

278

Yeah, I thought about that, the way in this diff seems very error-prone.

One solution is to modify the ProfileMap at the end (right before it call the write), this way is less intrusive. The trade-off is it could slow down the tool. Since we use the StringRef to save string copy, if we add brackets at the end, we have to create new string instance.

But since our tool is as fast, I‘m inclined to use a less intrusive way. what do you think?

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

278

Agreed, a final wrapping-up pass sounds solid. It might hurt performance a little bit but I'm not expecting a lot since each context string is only reprocessed once.

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

278

Cool, let me fix it!

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

278

+1 for final clean up pass. was thinking along the lines of that too.

[wlei]

use a final clean up pass for the bracket adding

[hoy]

In D95929#2542885, @wlei wrote:

use a final clean up pass for the bracket adding

Thanks! Looks a lot cleaner now.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

416

This should not be needed either. Perhaps in the future we'll also have the profile writer automatically detect pseudo-probe-based profiles and set FunctionSamples::ProfileIsProbeBased.

llvm/tools/llvm-profgen/ProfileGenerator.h

68–69

This should be removed since now every context has brackets.

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

416

How about in SampleProfWriter we use the checksum field to detect it's a probe-based profile then set the FunctionSamples::ProfileIsProbeBased?

llvm/tools/llvm-profgen/ProfileGenerator.h

68–69

It seems it also depends on this variable, see the code below. I saw your previous patch set it in the SampleProfReader but llvm-profgen only use the SampleProfWriter.

/// Return function name with context.
StringRef getNameWithContext(bool WithBracket = false) const {
  return FunctionSamples::ProfileIsCS
             ? Context.getNameWithContext(WithBracket)
             : Name;
}

I can understand we should automatically detect this, so how about to change the Writer to add some detection, say if the name has the bracket then set the ProfileIsCS?

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.h

68–69

Seems there aren't a central place in Writer to detect it, the write function is override by different format.
or how about code like this, we ensure the Context is always valid for CS profile, right?

StringRef getNameWithContext(bool WithBracket = false) const {
  return Context.getNameWithContext().size()?
          Context.getNameWithContext(WithBracket) : Name ;

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.h

68–69

Yeah, we might end up doing the detection per writer. That should only be needed for ProfileIsPseudoProbe . For now we don't need to set ProfileIsCS because the contexts all have brackets and they can just be treated as regular function profiles by the existing writers.

[wlei]

address Hongtao's feedback: remove ProfileIsCS assignment.

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.h

68–69

I see, I guess you meant ProfileIsCS is false by default, so it will use the default Name variable like the regular function and the Name variable actually has the brackets.

Just fixed following this idea by setting the Name.

FProfile.setName(FContext.getNameWithContext(true));

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.h

68–69

Exactly. Also we can avoid setting ProfileIsCS in llvm-profgen now. It's yet ready to do the same for ProfileIsProbeBased since work will be needed on the profile writer side.

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.h

68–69

Yeah, have removed the ProfileIsCS, see the latest diff.

[hoy]

llvm/test/tools/llvm-profgen/recursion-compression-pseudoprobe.test

10–11

What causes this diff? Looks like the output order changed.

llvm/tools/llvm-profgen/ProfileGenerator.h

68–69

Thanks!

[wlei]

llvm/test/tools/llvm-profgen/recursion-compression-pseudoprobe.test

10–11

Exactly, because now it ordered in the context with bracket, then the ascii code of ']' is greater than other alphabet or ':', the following test changes are the same.

[hoy]

LGTM.

[wenlei]

thanks for the cleanup, lgtm.

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

392

This seems problematic, the string could get destructed while we still keep the StringRef for it in FunctionSamples's name and context..

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

392

Does the try_emplace below copy construct a new one?

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

392

CxtWithBracketPMap is a StringMap, so key is StringRef not std::string.

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

392

Oh right, this is a problem. Good catch! How did you find it out?

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

392

I did think this is a problematic but after I used "Ret.first->first()" to get the StringRef, it seems something extend the lifetime of the string, otherwise Writer->write(CxtWithBracketPMap); will never work. do you have any case to trigger this?

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

392

I was doing some refactoring and a seemingly NFC change tripped this bug. I will send up a fix together with the refactoring.

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

392

The std::string object will be deallocated out once it goes out of the loop. But the underlying memory may not be overwritten depending on the heap usage in the system. It worked probably because no other memory allocation reclaims that piece of memory.

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

392

Thanks for the explanation. If so, the global ProfileMap is also problematic, IIRC, it also used a local std::string used similar as here..

At that time, I tried to avoid to do memory copy for the string, so didn't copy to a global lifetime container.

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

392

I think you are right. We are missing a global lifetime container. Probably ProfileMap should use std::string as keys.

llvm/lib/ProfileData/SampleProfWriter.cpp

 ///
 /// Note: it may be tempting to implement this in terms of
 /// FunctionSamples::print().  Please don't.  The dump functionality is intended
 /// for debugging and has no specified form.
 ///
 /// The format used here is more structured and deliberate because
 /// it needs to be parsed by the SampleProfileReaderText class.
 std::error_code SampleProfileWriterText::writeSample(const FunctionSamples &S) {
   auto &OS = *OutputStream;
-  if (FunctionSamples::ProfileIsCS)
-    OS << "[" << S.getNameWithContext() << "]:" << S.getTotalSamples();
-  else
-    OS << S.getName() << ":" << S.getTotalSamples();
+  OS << S.getNameWithContext(true) << ":" << S.getTotalSamples();
   if (Indent == 0)
     OS << ":" << S.getHeadSamples();
   OS << "\n";
 
   SampleSorter<LineLocation, SampleRecord> SortedSamples(S.getBodySamples());
   for (const auto &I : SortedSamples.get()) {
     LineLocation Loc = I->first;
     const SampleRecord &Sample = I->second;

llvm/test/tools/llvm-profgen/cs-extbinary.test

+; test for dwarf-based cs profile
+; RUN: llvm-profgen --format=extbinary --perfscript=%S/Inputs/recursion-compression-noprobe.perfscript --binary=%S/Inputs/recursion-compression-noprobe.perfbin --output=%t1 --csprof-cold-thres=0
+; RUN: llvm-profdata merge --sample --text --output=%t2 %t1
+; RUN: FileCheck %S/recursion-compression-noprobe.test --input-file %t2
+; RUN: llvm-profdata merge --sample --extbinary --output=%t3 %t2 && llvm-profdata merge --sample --text --output=%t4 %t3
+; RUN: diff -b %t2 %t4
+
+
+; test for probe-based cs profile
+; RUN: llvm-profgen --format=extbinary --perfscript=%S/Inputs/recursion-compression-pseudoprobe.perfscript --binary=%S/Inputs/recursion-compression-pseudoprobe.perfbin --output=%t5 --csprof-cold-thres=0
+; RUN: llvm-profdata merge --sample --text --output=%t6 %t5
+; RUN: FileCheck %S/recursion-compression-pseudoprobe.test --input-file %t6
+; RUN: llvm-profdata merge --sample --extbinary --output=%t7 %t6 && llvm-profdata merge --sample --text --output=%t8 %t7
+; RUN: diff -b %t6 %t8

llvm/test/tools/llvm-profgen/inline-cs-noprobe.test

 ; RUN: llvm-profgen --perfscript=%S/Inputs/inline-cs-noprobe.perfscript --binary=%S/Inputs/inline-cs-noprobe.perfbin --output=%t --show-unwinder-output --csprof-cold-thres=0 | FileCheck %s --check-prefix=CHECK-UNWINDER
 ; RUN: FileCheck %s --input-file %t
 
 ; CHECK:[main:1 @ foo]:44:0
-; CHECK: 2.2: 14
+; CHECK: 2.1: 14
 ; CHECK: 3: 15
-; CHECK: 3.2: 14 bar:14
-; CHECK: 3.4: 1
-; CHECK:[main:1 @ foo:3.2 @ bar]:14:0
+; CHECK: 3.1: 14 bar:14
+; CHECK: 3.2: 1
+; CHECK:[main:1 @ foo:3.1 @ bar]:14:0
 ; CHECK: 1: 14
 
 ; CHECK-UNWINDER: Binary(inline-cs-noprobe.perfbin)'s Range Counter:
 ; CHECK-UNWINDER: main:1 @ foo
 ; CHECK-UNWINDER:   (670, 6ad): 1
 ; CHECK-UNWINDER:   (67e, 69b): 1
 ; CHECK-UNWINDER:   (67e, 6ad): 13
 ; CHECK-UNWINDER:   (6bd, 6c8): 14
-; CHECK-UNWINDER: main:1 @ foo:3.2 @ bar
+; CHECK-UNWINDER: main:1 @ foo:3.1 @ bar
 ; CHECK-UNWINDER:   (6af, 6bb): 14
 
-
 ; CHECK-UNWINDER: Binary(inline-cs-noprobe.perfbin)'s Branch Counter:
 ; CHECK-UNWINDER: main:1 @ foo
 ; CHECK-UNWINDER:   (69b, 670): 1
 ; CHECK-UNWINDER:   (6c8, 67e): 15
 
 ; original code:
 ; clang -O3 -g test.c -o a.out
 #include <stdio.h>

llvm/test/tools/llvm-profgen/noinline-cs-noprobe.test

 ; CHECK-UNWINDER:   (5c8, 5dc): 2
 ; CHECK-UNWINDER:   (5d7, 5e5): 2
 ; CHECK-UNWINDER:   (5e9, 634): 3
 
 
 
 
 
+
+
 ; original code:
 ; clang -O0 -g test.c -o a.out
 #include <stdio.h>
 
 int bar(int x, int y) {
   if (x % 3) {
     return x - y;
   }

llvm/test/tools/llvm-profgen/recursion-compression-noprobe.test

 ; Firstly test uncompression(--compress-recursion=0)
 ; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression-noprobe.perfscript --binary=%S/Inputs/recursion-compression-noprobe.perfbin --output=%t --compress-recursion=0 --csprof-cold-thres=0
 ; RUN: FileCheck %s --input-file %t -check-prefix=CHECK-UNCOMPRESS
 ; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression-noprobe.perfscript --binary=%S/Inputs/recursion-compression-noprobe.perfbin --output=%t --csprof-cold-thres=0
 ; RUN: FileCheck %s --input-file %t
 
 ; CHECK-UNCOMPRESS:[main:1 @ foo:3 @ fa]:14:0
 ; CHECK-UNCOMPRESS: 1: 1
 ; CHECK-UNCOMPRESS: 2: 13 fb:11
 ; CHECK-UNCOMPRESS:[main:1 @ foo:3 @ fa:2 @ fb]:12:0
 ; CHECK-UNCOMPRESS: 1: 11
 ; CHECK-UNCOMPRESS: 2: 1 fa:1
-; CHECK-UNCOMPRESS:[main:1 @ foo:3 @ fa:2 @ fb:2 @ fa]:3:0
-; CHECK-UNCOMPRESS: 1: 1
-; CHECK-UNCOMPRESS: 2: 2 fb:1
 ; CHECK-UNCOMPRESS:[main:1 @ foo]:3:0
 ; CHECK-UNCOMPRESS: 2: 1
 ; CHECK-UNCOMPRESS: 3: 2 fa:1
-; CHECK-UNCOMPRESS:[main:1 @ foo:3 @ fa:2 @ fb:2 @ fa:2 @ fb:2 @ fa]:1:0
-; CHECK-UNCOMPRESS: 4: 1
+; CHECK-UNCOMPRESS:[main:1 @ foo:3 @ fa:2 @ fb:2 @ fa]:3:0
+; CHECK-UNCOMPRESS: 1: 1
+; CHECK-UNCOMPRESS: 2: 2 fb:1
 ; CHECK-UNCOMPRESS:[main:1 @ foo:3 @ fa:2 @ fb:2 @ fa:2 @ fb]:1:0
 ; CHECK-UNCOMPRESS: 2: 1 fa:1
+; CHECK-UNCOMPRESS:[main:1 @ foo:3 @ fa:2 @ fb:2 @ fa:2 @ fb:2 @ fa]:1:0
+; CHECK-UNCOMPRESS: 4: 1
+
 
 ; CHECK: [main:1 @ foo:3 @ fa]:14:0
 ; CHECK:  1: 1
 ; CHECK:  2: 13 fb:11
 ; CHECK: [main:1 @ foo:3 @ fa:2 @ fb]:12:0
 ; CHECK:  1: 11
 ; CHECK:  2: 1 fa:1
 ; CHECK: [main:1 @ foo:3 @ fa:2 @ fb:2 @ fa]:4:0

llvm/test/tools/llvm-profgen/recursion-compression-pseudoprobe.test

 ; Firstly test uncompression(--compress-recursion=0)
 ; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression-pseudoprobe.perfscript --binary=%S/Inputs/recursion-compression-pseudoprobe.perfbin --output=%t --compress-recursion=0 --csprof-cold-thres=0
 ; RUN: FileCheck %s --input-file %t -check-prefix=CHECK-UNCOMPRESS
 ; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression-pseudoprobe.perfscript --binary=%S/Inputs/recursion-compression-pseudoprobe.perfbin --output=%t --show-unwinder-output --csprof-cold-thres=0 | FileCheck %s --check-prefix=CHECK-UNWINDER
 ; RUN: FileCheck %s --input-file %t
 
-; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa]:4:1
+; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa]:4:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  3: 1
-; CHECK-UNCOMPRESS:  4: 1
-; CHECK-UNCOMPRESS:  7: 1 fb:1
+; CHECK-UNCOMPRESS:  5: 1
+; CHECK-UNCOMPRESS:  8: 1 fa:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 120515930909
 ; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa]:4:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  3: 1
 ; CHECK-UNCOMPRESS:  4: 1
 ; CHECK-UNCOMPRESS:  7: 1 fb:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 120515930909
-; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa]:4:1
+; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa]:4:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  3: 1
-; CHECK-UNCOMPRESS:  5: 1
-; CHECK-UNCOMPRESS:  8: 1 fa:1
+; CHECK-UNCOMPRESS:  4: 1
+; CHECK-UNCOMPRESS:  7: 1 fb:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 120515930909
-; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa:7 @ fb]:3:1
-; CHECK-UNCOMPRESS:  1: 1
-; CHECK-UNCOMPRESS:  3: 1
-; CHECK-UNCOMPRESS:  6: 1 fa:1
-; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
-; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb]:3:1
-; CHECK-UNCOMPRESS:  1: 1
-; CHECK-UNCOMPRESS:  3: 1
-; CHECK-UNCOMPRESS:  6: 1 fa:1
-; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
-; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
-; CHECK-UNCOMPRESS:  1: 1
-; CHECK-UNCOMPRESS:  3: 1
-; CHECK-UNCOMPRESS:  6: 1 fa:1
-; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
-; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
+; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  2: 1
 ; CHECK-UNCOMPRESS:  5: 1 fb:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
 ; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  2: 1
 ; CHECK-UNCOMPRESS:  5: 1 fb:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
-; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
+; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  2: 1
 ; CHECK-UNCOMPRESS:  5: 1 fb:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
+; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
+; CHECK-UNCOMPRESS:  1: 1
+; CHECK-UNCOMPRESS:  3: 1
+; CHECK-UNCOMPRESS:  6: 1 fa:1
+; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
+; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb]:3:1
+; CHECK-UNCOMPRESS:  1: 1
+; CHECK-UNCOMPRESS:  3: 1
+; CHECK-UNCOMPRESS:  6: 1 fa:1
+; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
+; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa:7 @ fb]:3:1
+; CHECK-UNCOMPRESS:  1: 1
+; CHECK-UNCOMPRESS:  3: 1
+; CHECK-UNCOMPRESS:  6: 1 fa:1
+; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
 ; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa:7 @ fb:6 @ fa]:2:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  3: 1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 120515930909
 ; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:1:0
 ; CHECK-UNCOMPRESS:  5: 1 fb:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
 
 
 ; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb]:13:4
 ; CHECK:  1: 4
 ; CHECK:  2: 3
 ; CHECK:  3: 1
 ; CEHCK:  5: 4 fb:4
 ; CHECK:  6: 1 fa:1
 ; CHECK !CFGChecksum: 72617220756
 ; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa]:6:2
 ; CHECK:  1: 2
 ; CHECK:  3: 2
 ; CHECK:  4: 1
 ; CHECK:  7: 1 fb:1
 ; CHECK:  !CFGChecksum: 120515930909
-; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa]:4:1
-; CHECK:  1: 1
-; CHECK:  3: 1
-; CHECK:  4: 1
-; CHECK:  7: 1 fb:1
-; CHECK:  !CFGChecksum: 120515930909
 ; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa]:4:1
 ; CHECK:  1: 1
 ; CHECK:  3: 1
 ; CHECK:  5: 1
 ; CHECK:  8: 1 fa:1
 ; CHECK:  !CFGChecksum: 120515930909
-; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa:7 @ fb]:3:1
+; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa]:4:1
+; CHECK:  1: 1
+; CHECK:  3: 1
+; CHECK:  4: 1
+; CHECK:  7: 1 fb:1
+; CHECK:  !CFGChecksum: 120515930909
+; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb]:3:1
 ; CHECK:  1: 1
 ; CHECK:  3: 1
 ; CHECK:  6: 1 fa:1
 ; CHECK:  !CFGChecksum: 72617220756
-; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb]:3:1
+; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa:7 @ fb]:3:1
 ; CHECK:  1: 1
 ; CHECK:  3: 1
 ; CHECK:  6: 1 fa:1
 ; CHECK:  !CFGChecksum: 72617220756
 
 
+
 ; CHECK-UNWINDER: Binary(recursion-compression-pseudoprobe.perfbin)'s Range Counter:
 ; CHECK-UNWINDER: main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5
 ; CHECK-UNWINDER:   (7a0, 7a7): 1
 ; CHECK-UNWINDER:   (7a0, 7ab): 3
 ; CHECK-UNWINDER:   (7b2, 7b5): 1
 ; CHECK-UNWINDER: main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6
 ; CHECK-UNWINDER:   (7c0, 7d4): 1
 ; CHECK-UNWINDER: main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8

llvm/tools/llvm-profgen/ProfileGenerator.h

 
 public:
   ProfileGenerator(){};
   virtual ~ProfileGenerator() = default;
   static std::unique_ptr<ProfileGenerator>
   create(const BinarySampleCounterMap &BinarySampleCounters,
          enum PerfScriptType SampleType);
   virtual void generateProfile() = 0;
-  // Merge and trim profile with cold context before serialization,
-  // only eligible for CS profile
-  virtual void
-  mergeAndTrimColdProfile(StringMap<FunctionSamples> &ProfileMap){};
   // Use SampleProfileWriter to serialize profile map
+  virtual void write(std::unique_ptr<SampleProfileWriter> Writer,
+                     StringMap<FunctionSamples> &ProfileMap);
   void write();
 
 protected:
   /*
   For each region boundary point, mark if it is begin or end (or both) of
   the region. Boundary points are inclusive. Log the sample count as well
   so we can use it when we compute the sample count of each disjoint region
   later. Note that there might be multiple ranges with different sample

 protected:
   const BinarySampleCounterMap &BinarySampleCounters;
 
 public:
   CSProfileGenerator(const BinarySampleCounterMap &Counters)
       : BinarySampleCounters(Counters){};
 
 public:
   void generateProfile() override {
-    // Enable context-sensitive functionalities in SampleProf
-    FunctionSamples::ProfileIsCS = true;
     for (const auto &BI : BinarySampleCounters) {
       ProfiledBinary *Binary = BI.first;
       for (const auto &CI : BI.second) {
         const StringBasedCtxKey *CtxKey =
             dyn_cast<StringBasedCtxKey>(CI.first.getPtr());
         StringRef ContextId(CtxKey->Context);
         // Get or create function profile for the range
         FunctionSamples &FunctionProfile =
             getFunctionProfileForContext(ContextId);
static void compressRecursionContext(SmallVectorImpl<T> &Context,
     }
   }
 
 protected:
   // Lookup or create FunctionSamples for the context
   FunctionSamples &getFunctionProfileForContext(StringRef ContextId);
   // Merge cold context profile whose total sample is below threshold
   // into base profile.
-  void mergeAndTrimColdProfile(StringMap<FunctionSamples> &ProfileMap) override;
+  void mergeAndTrimColdProfile(StringMap<FunctionSamples> &ProfileMap);
+  void write(std::unique_ptr<SampleProfileWriter> Writer,
+             StringMap<FunctionSamples> &ProfileMap) override;
 
 private:
   // Helper function for updating body sample for a leaf location in
   // FunctionProfile
   void updateBodySamplesforFunctionProfile(FunctionSamples &FunctionProfile,
                                            const FrameLocation &LeafLoc,
                                            uint64_t Count);
   void populateFunctionBodySamples(FunctionSamples &FunctionProfile,

llvm/tools/llvm-profgen/ProfileGenerator.cpp

ProfileGenerator::create(const BinarySampleCounterMap &BinarySampleCounters,
   } else {
     // TODO:
     llvm_unreachable("Unsupported perfscript!");
   }
 
   return ProfileGenerator;
 }
 
+void ProfileGenerator::write(std::unique_ptr<SampleProfileWriter> Writer,
+                             StringMap<FunctionSamples> &ProfileMap) {
+  Writer->write(ProfileMap);
+}
+
 void ProfileGenerator::write() {
   auto WriterOrErr = SampleProfileWriter::create(OutputFilename, OutputFormat);
   if (std::error_code EC = WriterOrErr.getError())
     exitWithError(EC, OutputFilename);
-  auto Writer = std::move(WriterOrErr.get());
-  mergeAndTrimColdProfile(ProfileMap);
-  Writer->write(ProfileMap);
+  write(std::move(WriterOrErr.get()), ProfileMap);
 }
 
 void ProfileGenerator::findDisjointRanges(RangeSample &DisjointRanges,
                                           const RangeSample &Ranges) {
 
   /*
   Regions may overlap with each other. Using the boundary info, find all
   disjoint ranges and their sample count. BoundaryPoint contains the count

 }
 
 FunctionSamples &
 CSProfileGenerator::getFunctionProfileForContext(StringRef ContextStr) {
   auto Ret = ProfileMap.try_emplace(ContextStr, FunctionSamples());
   if (Ret.second) {
     SampleContext FContext(Ret.first->first(), RawContext);
     FunctionSamples &FProfile = Ret.first->second;
-    FProfile.setName(FContext.getNameWithoutContext());
     FProfile.setContext(FContext);
   }
   return Ret.first->second;
 }
 
 void CSProfileGenerator::updateBodySamplesforFunctionProfile(
     FunctionSamples &FunctionProfile, const FrameLocation &LeafLoc,
     uint64_t Count) {
for (auto Entry : BranchCounters) {
 
     // Record called target sample and its count
     const FrameLocation &LeafLoc = Binary->getInlineLeafFrameLoc(SourceOffset);
     FunctionProfile.addCalledTargetSamples(LeafLoc.second.LineOffset,
                                            LeafLoc.second.Discriminator,
                                            CalleeName, Count);
 
     // Record head sample for called target(callee)
-    // TODO: Cleanup ' @ '
-    std::string CalleeContextId =
-        getCallSite(LeafLoc) + " @ " + CalleeName.str();
+    std::ostringstream OCalleeCtxStr;
     if (ContextId.find(" @ ") != StringRef::npos) {
-      CalleeContextId =
-          ContextId.rsplit(" @ ").first.str() + " @ " + CalleeContextId;
+      OCalleeCtxStr << ContextId.rsplit(" @ ").first.str();
+      OCalleeCtxStr << " @ ";
     }
+    OCalleeCtxStr << getCallSite(LeafLoc) << " @ " << CalleeName.str();
 
     FunctionSamples &CalleeProfile =
-        getFunctionProfileForContext(CalleeContextId);
+        getFunctionProfileForContext(OCalleeCtxStr.str());
     assert(Count != 0 && "Unexpected zero weight branch");
     CalleeProfile.addHeadSamples(Count);
   }
 }
 
 static FrameLocation getCallerContext(StringRef CalleeContext,
                                       StringRef &CallerNameWithContext) {
   StringRef CallerContext = CalleeContext.rsplit(" @ ").first;
for (const auto &Item : ProfileMap) {
     // Since we don't have call count for inlined functions, we
     // estimate it from inlinee's profile using entry body sample.
     uint64_t EstimatedCallCount = CalleeProfile.getEntrySamples();
     // If we don't have samples with location, use 1 to indicate live.
     if (!EstimatedCallCount && !CalleeProfile.getBodySamples().size())
       EstimatedCallCount = 1;
     CallerProfile.addCalledTargetSamples(
         CallerLeafFrameLoc.second.LineOffset,
-        CallerLeafFrameLoc.second.Discriminator, CalleeProfile.getName(),
-        EstimatedCallCount);
+        CallerLeafFrameLoc.second.Discriminator,
+        CalleeProfile.getContext().getNameWithoutContext(), EstimatedCallCount);
     CallerProfile.addBodySamples(CallerLeafFrameLoc.second.LineOffset,
                                  CallerLeafFrameLoc.second.Discriminator,
                                  EstimatedCallCount);
     CallerProfile.addTotalSamples(EstimatedCallCount);
   }
 }
 
 void CSProfileGenerator::mergeAndTrimColdProfile(
for (const auto &I : ProfileMap) {
     if (FunctionProfile.getTotalSamples() >= CSProfColdThres)
       continue;
     ToRemoveVec.emplace_back(I.getKey(), &I.second);
   }
 
   // Remove the code profile from ProfileMap and merge them into BaseProileMap
   StringMap<FunctionSamples> BaseProfileMap;
   for (const auto &I : ToRemoveVec) {
-    auto Ret =
-        BaseProfileMap.try_emplace(I.second->getName(), FunctionSamples());
+    auto Ret = BaseProfileMap.try_emplace(
+        I.second->getContext().getNameWithoutContext(), FunctionSamples());
     FunctionSamples &BaseProfile = Ret.first->second;
     BaseProfile.merge(*I.second);
     ProfileMap.erase(I.first);
   }
 
   // Merge the base profiles into ProfileMap;
   for (const auto &I : BaseProfileMap) {
     // Filter the cold base profile
     if (!CSProfKeepCold && I.second.getTotalSamples() < CSProfColdThres &&
         ProfileMap.find(I.getKey()) == ProfileMap.end())
       continue;
     // Merge the profile if the original profile exists, otherwise just insert
     // as a new profile
     FunctionSamples &OrigProfile = getFunctionProfileForContext(I.getKey());
-    StringRef TmpName = OrigProfile.getName();
     OrigProfile.merge(I.second);
-    // Should use the name ref from ProfileMap's key to avoid name being freed
-    // from BaseProfileMap
-    OrigProfile.setName(TmpName);
   }
 }
 
+void CSProfileGenerator::write(std::unique_ptr<SampleProfileWriter> Writer,
+                               StringMap<FunctionSamples> &ProfileMap) {
+  mergeAndTrimColdProfile(ProfileMap);
+  // Add bracket for context key to support different profile binary format
+  StringMap<FunctionSamples> CxtWithBracketPMap;
+  for (const auto &Item : ProfileMap) {
+    std::string ContextWithBracket = "[" + Item.first().str() + "]";
+    auto Ret = CxtWithBracketPMap.try_emplace(ContextWithBracket, Item.second);
+    assert(Ret.second && "Must be a unique context");
+    SampleContext FContext(Ret.first->first(), RawContext);
+    FunctionSamples &FProfile = Ret.first->second;
+    FProfile.setName(FContext.getNameWithContext(true));
+    FProfile.setContext(FContext);
+  }
+  Writer->write(CxtWithBracketPMap);
+}
+
 // Helper function to extract context prefix string stack
 // Extract context stack for reusing, leaf context stack will
 // be added compressed while looking up function profile
 static void
 extractPrefixContextStack(SmallVectorImpl<std::string> &ContextStrStack,
                           const SmallVectorImpl<const PseudoProbe *> &Probes,
                           ProfiledBinary *Binary) {
   for (const auto *P : Probes) {
     Binary->getInlineContextForProbe(P, ContextStrStack, true);
   }
 }
 
 void PseudoProbeCSProfileGenerator::generateProfile() {
-  // Enable CS and pseudo probe functionalities in SampleProf
-  FunctionSamples::ProfileIsCS = true;
+  // Enable pseudo probe functionalities in SampleProf
   FunctionSamples::ProfileIsProbeBased = true;
   for (const auto &BI : BinarySampleCounters) {
     ProfiledBinary *Binary = BI.first;
     for (const auto &CI : BI.second) {
       const ProbeBasedCtxKey *CtxKey =
           dyn_cast<ProbeBasedCtxKey>(CI.first.getPtr());
       SmallVector<std::string, 16> ContextStrStack;
       extractPrefixContextStack(ContextStrStack, CtxKey->Probes, Binary);
if (Probe->isEntry()) {
         uint64_t CallerIndex = CallerLeafFrameLoc.second.LineOffset;
         assert(CallerIndex &&
                "Inferred caller's location index shouldn't be zero!");
         FunctionSamples &CallerProfile =
             getFunctionProfileForContext(CallerContextId);
         CallerProfile.setFunctionHash(InlinerDesc->FuncHash);
         CallerProfile.addBodySamples(CallerIndex, 0, Count);
         CallerProfile.addTotalSamples(Count);
-        CallerProfile.addCalledTargetSamples(CallerIndex, 0,
-                                             FunctionProfile.getName(), Count);
+        CallerProfile.addCalledTargetSamples(
+            CallerIndex, 0,
+            FunctionProfile.getContext().getNameWithoutContext(), Count);
       }
     }
   }
 }
 
 void PseudoProbeCSProfileGenerator::populateBoundarySamplesWithProbes(
     const BranchSample &BranchCounter,
     SmallVectorImpl<std::string> &ContextStrStack, ProfiledBinary *Binary) {

llvm/tools/llvm-profgen/ProfiledBinary.cpp

 //===-- ProfiledBinary.cpp - Binary decoder ---------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "ProfiledBinary.h"
 #include "ErrorHandling.h"
 #include "ProfileGenerator.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Demangle/Demangle.h"
+#include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Format.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/TargetSelect.h"
 
 #define DEBUG_TYPE "load-binary"
 
 using namespace llvm;
FrameLocationStack ProfiledBinary::symbolize(const InstructionPointer &IP,
   for (int32_t I = InlineStack.getNumberOfFrames() - 1; I >= 0; I--) {
     const auto &CallerFrame = InlineStack.getFrame(I);
     if (CallerFrame.FunctionName == "<invalid>")
       break;
     StringRef FunctionName(CallerFrame.FunctionName);
     if (UseCanonicalFnName)
       FunctionName = FunctionSamples::getCanonicalFnName(FunctionName);
     LineLocation Line(CallerFrame.Line - CallerFrame.StartLine,
-                      CallerFrame.Discriminator);
+                      DILocation::getBaseDiscriminatorFromDiscriminator(
+                          CallerFrame.Discriminator));
     FrameLocation Callsite(FunctionName.str(), Line);
     CallStack.push_back(Callsite);
   }
 
   return CallStack;
 }
 
 InstructionPointer::InstructionPointer(ProfiledBinary *Binary, uint64_t Address,