Index: include/llvm/Passes/PassBuilder.h
===================================================================
--- include/llvm/Passes/PassBuilder.h
+++ include/llvm/Passes/PassBuilder.h
@@ -32,6 +32,7 @@
   std::string ProfileGenFile = "";
   std::string ProfileUseFile = "";
   std::string SampleProfileFile = "";
+  bool SampleProfileGen = false;
   bool RunProfileGen = false;
 };
 
Index: lib/Passes/PassBuilder.cpp
===================================================================
--- lib/Passes/PassBuilder.cpp
+++ lib/Passes/PassBuilder.cpp
@@ -330,6 +330,10 @@
   assert(Level != O0 && "Must request optimizations!");
   FunctionPassManager FPM(DebugLogging);
 
+  // If SamplePGO is enabled, add the AddDiscriminator pass.
+  if (PGOOpt && (PGOpt.SampleProfileGen || PGOOpt.SampleProfileFile))
+    FPM.addPass(AddDiscriminatorsPass());
+
   // Form SSA out of local memory accesses after breaking apart aggregates into
   // scalars.
   FPM.addPass(SROA());