[llvm] [PGO] Add uniformity profile metadata and gate CHR (PR #221517)

[Matt Arsenault via llvm-commits]

================
@@ -1788,29 +1788,40 @@ void PGOUseFunc::setBlockUniformityAttribute() {
   if (ProfileRecord.UniformityBits.empty())
     return;
 
-  // Annotate uniformity on each instrumented IR basic block so later codegen
-  // passes (MachineFunction) can consume it without relying on fragile block
-  // numbering heuristics.
-  //
-  // Metadata kind: LLVMContext::MD_block_uniformity_profile
-  // Payload: i1 (true = uniform, false = divergent)
+  // Mark the function as having uniformity profile, then annotate each uniform
+  // instrumented IR basic block so later codegen passes (MachineFunction) can
+  // consume it without relying on fragile block numbering heuristics.
+  // Metadata presence on a terminator means uniform; divergent blocks have no
+  // terminator metadata.
 
   std::vector<BasicBlock *> InstrumentBBs;
   FuncInfo.getInstrumentBBs(InstrumentBBs);
 
   LLVMContext &Ctx = F.getContext();
-  Type *Int1Ty = Type::getInt1Ty(Ctx);
-
+  MDNode *UniformMD = MDNode::get(Ctx, {});
+  F.setMetadata(LLVMContext::MD_block_uniformity_profile, UniformMD);
+  DenseMap<CondBrInst *, bool> BranchUniformity;
   for (size_t I = 0, E = InstrumentBBs.size(); I < E; ++I) {
     BasicBlock *BB = InstrumentBBs[I];
     if (!BB || !BB->getTerminator())
       continue;
     bool IsUniform = ProfileRecord.isBlockUniform(I);
-    auto *MD = MDNode::get(
-        Ctx, ConstantAsMetadata::get(ConstantInt::get(Int1Ty, IsUniform)));
-    BB->getTerminator()->setMetadata(LLVMContext::MD_block_uniformity_profile,
-                                     MD);
+    // A counter placed in a block with a single conditional predecessor also
+    // measures the active lanes on that outgoing edge. Record the branch as
+    // uniform only when every instrumented outgoing edge is uniform.
+    if (BasicBlock *Pred = BB->getSinglePredecessor()) {
+      if (auto *Branch = dyn_cast<CondBrInst>(Pred->getTerminator())) {
+        auto It = BranchUniformity.try_emplace(Branch, true).first;
+        It->second &= IsUniform;
+      }
+    }
+    if (IsUniform)
+      BB->getTerminator()->setMetadata(LLVMContext::MD_block_uniformity_profile,
+                                       UniformMD);
   }
+  for (auto [Branch, IsUniform] : BranchUniformity)
----------------
arsenm wrote:

Braces 

https://github.com/llvm/llvm-project/pull/221517