[llvm] [CodeGen] Add OffloadBlockUniformityAnalysis for offload PGO (PR #178417)

Yaxun Liu via llvm-commits llvm-commits at lists.llvm.org
Thu Jul 2 19:57:16 PDT 2026


https://github.com/yxsamliu updated https://github.com/llvm/llvm-project/pull/178417

>From caca6c7ed1a00c36f13a203f03dcbc691f166b7d Mon Sep 17 00:00:00 2001
From: "Yaxun (Sam) Liu" <yaxun.liu at amd.com>
Date: Thu, 5 Feb 2026 11:52:20 -0500
Subject: [PATCH 1/2] [CodeGen] Add BlockUniformityProfile for PGO-guided spill
 placement

Add BlockUniformityProfile analysis that reads block uniformity metadata
from PGO profiles to guide register allocation decisions.

When PGO data indicates a basic block executes divergently, SpillPlacement
flattens its frequency to the function entry frequency. This prevents PGO
from biasing spill decisions toward blocks that appear hot in profiles but
run with partial wave occupancy.

Key components:
- PGOInstrumentation emits block.uniformity.profile metadata during PGO use,
  keyed by the fixed metadata kind MD_block_uniformity_profile.
- BlockUniformityProfile reads !block.uniformity.profile metadata.
- BlockUniformityProfileProxy provides the analysis for the new pass manager.
- BlockUniformityProfilePrinterPass provides a debug printer pass.
- SpillPlacement flattens divergent block frequencies.

The metadata is attached during PGO use by PGOInstrumentation.cpp and uses
a fixed metadata kind for efficient lookup.
---
 .../llvm/CodeGen/BlockUniformityProfile.h     |  79 ++++++++++++
 llvm/include/llvm/CodeGen/SpillPlacement.h    |   4 +-
 llvm/include/llvm/Passes/CodeGenPassBuilder.h |   1 +
 .../llvm/Passes/MachinePassRegistry.def       |   4 +
 llvm/lib/CodeGen/BlockUniformityProfile.cpp   | 121 ++++++++++++++++++
 llvm/lib/CodeGen/CMakeLists.txt               |   1 +
 llvm/lib/CodeGen/SpillPlacement.cpp           |  24 +++-
 llvm/lib/Passes/PassBuilder.cpp               |   1 +
 .../AMDGPU/block-uniformity-profile.ll        |  79 ++++++++++++
 9 files changed, 308 insertions(+), 6 deletions(-)
 create mode 100644 llvm/include/llvm/CodeGen/BlockUniformityProfile.h
 create mode 100644 llvm/lib/CodeGen/BlockUniformityProfile.cpp
 create mode 100644 llvm/test/CodeGen/AMDGPU/block-uniformity-profile.ll

diff --git a/llvm/include/llvm/CodeGen/BlockUniformityProfile.h b/llvm/include/llvm/CodeGen/BlockUniformityProfile.h
new file mode 100644
index 0000000000000..209e0f893b6ca
--- /dev/null
+++ b/llvm/include/llvm/CodeGen/BlockUniformityProfile.h
@@ -0,0 +1,79 @@
+//===- BlockUniformityProfile.h - Block uniformity from PGO -*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Provide per-(Machine)basic-block uniformity information from PGO profiles.
+//
+// The source of truth is IR metadata attached during PGO use:
+//   - Metadata name: "block.uniformity.profile"
+//   - Payload: i1 (true = uniform, false = divergent)
+//
+// This is intentionally target-agnostic: any backend that produces
+// uniformity bits in the profile can attach the same metadata and reuse this
+// proxy in codegen.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_BLOCKUNIFORMITYPROFILE_H
+#define LLVM_CODEGEN_BLOCKUNIFORMITYPROFILE_H
+
+#include "llvm/ADT/BitVector.h"
+#include "llvm/CodeGen/MachineFunctionAnalysis.h"
+#include "llvm/CodeGen/MachineFunctionAnalysisManager.h"
+#include "llvm/CodeGen/MachinePassManager.h"
+#include "llvm/Support/Compiler.h"
+
+namespace llvm {
+
+class MachineBasicBlock;
+class MachineFunction;
+class raw_ostream;
+
+class BlockUniformityProfile {
+public:
+  LLVM_ABI void compute(const MachineFunction &MF);
+
+  bool hasProfile() const { return HasProfile; }
+
+  // Returns true if the block is considered divergent. If profile exists for
+  // the function but a block has no explicit annotation, it is treated as
+  // divergent (conservative).
+  LLVM_ABI bool isDivergent(const MachineBasicBlock &MBB) const;
+
+  LLVM_ABI void print(raw_ostream &OS, const MachineFunction &MF) const;
+
+private:
+  bool HasProfile = false;
+  unsigned NumBlockIDs = 0;
+  BitVector DivergentBlocks;
+};
+
+class BlockUniformityProfileProxy
+    : public AnalysisInfoMixin<BlockUniformityProfileProxy> {
+  friend AnalysisInfoMixin<BlockUniformityProfileProxy>;
+  static AnalysisKey Key;
+
+public:
+  using Result = BlockUniformityProfile;
+  LLVM_ABI Result run(MachineFunction &MF,
+                      MachineFunctionAnalysisManager &MFAM);
+};
+
+class BlockUniformityProfilePrinterPass
+    : public PassInfoMixin<BlockUniformityProfilePrinterPass> {
+  raw_ostream &OS;
+
+public:
+  explicit BlockUniformityProfilePrinterPass(raw_ostream &OS) : OS(OS) {}
+  LLVM_ABI PreservedAnalyses run(MachineFunction &MF,
+                                 MachineFunctionAnalysisManager &MFAM);
+  static bool isRequired() { return true; }
+};
+
+} // end namespace llvm
+
+#endif // LLVM_CODEGEN_BLOCKUNIFORMITYPROFILE_H
diff --git a/llvm/include/llvm/CodeGen/SpillPlacement.h b/llvm/include/llvm/CodeGen/SpillPlacement.h
index bc206075d9ba4..857aa37f0b301 100644
--- a/llvm/include/llvm/CodeGen/SpillPlacement.h
+++ b/llvm/include/llvm/CodeGen/SpillPlacement.h
@@ -36,6 +36,7 @@
 namespace llvm {
 
 class BitVector;
+class BlockUniformityProfile;
 class EdgeBundles;
 class MachineBlockFrequencyInfo;
 class MachineFunction;
@@ -169,7 +170,8 @@ class SpillPlacement {
   LLVM_ABI void releaseMemory();
 
   void run(MachineFunction &MF, EdgeBundles *Bundles,
-           MachineBlockFrequencyInfo *MBFI);
+           MachineBlockFrequencyInfo *MBFI,
+           const BlockUniformityProfile *Profile);
   void activate(unsigned n);
   void setThreshold(BlockFrequency Entry);
 
diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index 3840b90c70811..04aecd9495778 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -25,6 +25,7 @@
 #include "llvm/Analysis/TypeBasedAliasAnalysis.h"
 #include "llvm/CodeGen/AsmPrinter.h"
 #include "llvm/CodeGen/AsmPrinterAnalysis.h"
+#include "llvm/CodeGen/BlockUniformityProfile.h"
 #include "llvm/CodeGen/BranchFoldingPass.h"
 #include "llvm/CodeGen/CodeGenPrepare.h"
 #include "llvm/CodeGen/DeadMachineInstructionElim.h"
diff --git a/llvm/include/llvm/Passes/MachinePassRegistry.def b/llvm/include/llvm/Passes/MachinePassRegistry.def
index cb3aa9c013c43..bf6bca62e942e 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -26,6 +26,8 @@
 // LiveVariables can be removed completely, and LiveIntervals can be directly
 // computed. (We still either need to regenerate kill flags after regalloc, or
 // preferably fix the scavenger to not depend on them).
+MACHINE_FUNCTION_ANALYSIS("block-uniformity-profile",
+                          BlockUniformityProfileProxy())
 MACHINE_FUNCTION_ANALYSIS("edge-bundles", EdgeBundlesAnalysis())
 MACHINE_FUNCTION_ANALYSIS("gisel-cse-analysis", GISelCSEAnalysis(TM))
 MACHINE_FUNCTION_ANALYSIS("gisel-value-tracking", GISelValueTrackingAnalysis())
@@ -98,6 +100,8 @@ MACHINE_FUNCTION_PASS("postra-machine-sink", PostRAMachineSinkingPass())
 MACHINE_FUNCTION_PASS("postmisched", PostMachineSchedulerPass(TM))
 MACHINE_FUNCTION_PASS("post-ra-pseudos", ExpandPostRAPseudosPass())
 MACHINE_FUNCTION_PASS("print", PrintMIRPass())
+MACHINE_FUNCTION_PASS("print<block-uniformity-profile>",
+                      BlockUniformityProfilePrinterPass(errs()))
 MACHINE_FUNCTION_PASS("print<gisel-value-tracking>", GISelValueTrackingPrinterPass(errs()))
 MACHINE_FUNCTION_PASS("print<livedebugvars>", LiveDebugVariablesPrinterPass(errs()))
 MACHINE_FUNCTION_PASS("print<live-intervals>", LiveIntervalsPrinterPass(errs()))
diff --git a/llvm/lib/CodeGen/BlockUniformityProfile.cpp b/llvm/lib/CodeGen/BlockUniformityProfile.cpp
new file mode 100644
index 0000000000000..fff9843ce59eb
--- /dev/null
+++ b/llvm/lib/CodeGen/BlockUniformityProfile.cpp
@@ -0,0 +1,121 @@
+//===- BlockUniformityProfile.cpp - Block uniformity from PGO -----------===//
+//
+// 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 "llvm/CodeGen/BlockUniformityProfile.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/Support/raw_ostream.h"
+#include <optional>
+
+using namespace llvm;
+
+static std::optional<bool> getIRBlockUniformity(const BasicBlock &BB) {
+  const Instruction *TI = BB.getTerminator();
+  if (!TI)
+    return std::nullopt;
+
+  MDNode *MD = TI->getMetadata(LLVMContext::MD_block_uniformity_profile);
+  if (!MD)
+    return std::nullopt;
+
+  // Metadata format: !{i1 IsUniform} - structural validity assumed (verifier).
+  // Returns true if uniform (not divergent).
+  return mdconst::extract<ConstantInt>(MD->getOperand(0))->isOne();
+}
+
+void BlockUniformityProfile::compute(const MachineFunction &MF) {
+  HasProfile = false;
+  NumBlockIDs = MF.getNumBlockIDs();
+  DivergentBlocks.clear();
+  DivergentBlocks.resize(NumBlockIDs);
+
+  // First determine whether any uniformity profile exists for this function.
+  for (const MachineBasicBlock &MBB : MF) {
+    const BasicBlock *BB = MBB.getBasicBlock();
+    if (!BB)
+      continue;
+    if (getIRBlockUniformity(*BB).has_value()) {
+      HasProfile = true;
+      break;
+    }
+  }
+
+  if (!HasProfile)
+    return;
+
+  // Conservative behavior: if profile exists for the function but we
+  // cannot classify a particular (Machine)basic block, treat it as divergent.
+  for (const MachineBasicBlock &MBB : MF) {
+    const unsigned Num = MBB.getNumber();
+    bool IsDivergent = true;
+    if (const BasicBlock *BB = MBB.getBasicBlock()) {
+      if (auto U = getIRBlockUniformity(*BB))
+        IsDivergent = !*U; // Metadata stores IsUniform, we want IsDivergent
+    }
+    if (Num < DivergentBlocks.size() && IsDivergent)
+      DivergentBlocks.set(Num);
+  }
+}
+
+void BlockUniformityProfile::print(raw_ostream &OS,
+                                   const MachineFunction &MF) const {
+  OS << "BlockUniformityProfile for function: ";
+  MF.getFunction().printAsOperand(OS, /*PrintType=*/false);
+  OS << '\n';
+  OS << "HasProfile: " << (HasProfile ? "true" : "false") << '\n';
+  if (!HasProfile)
+    return;
+
+  for (const MachineBasicBlock &MBB : MF) {
+    const BasicBlock *BB = MBB.getBasicBlock();
+    if (!BB)
+      continue;
+    OS << "  " << printMBBReference(MBB);
+    if (BB->hasName())
+      OS << " (%" << BB->getName() << ")";
+    if (auto U = getIRBlockUniformity(*BB)) {
+      OS << ": " << (*U ? "uniform" : "divergent") << '\n';
+      continue;
+    }
+    OS << ": no PGO annotation (treated divergent for spill placement)\n";
+  }
+}
+
+bool BlockUniformityProfile::isDivergent(const MachineBasicBlock &MBB) const {
+  if (!HasProfile)
+    return false;
+  assert(MBB.getParent()->getNumBlockIDs() == NumBlockIDs &&
+         "MachineFunction was modified without invalidating "
+         "BlockUniformityProfile");
+  const unsigned Num = MBB.getNumber();
+  assert(Num < DivergentBlocks.size() && "Block number out of range");
+  return DivergentBlocks.test(Num);
+}
+
+AnalysisKey BlockUniformityProfileProxy::Key;
+
+BlockUniformityProfileProxy::Result
+BlockUniformityProfileProxy::run(MachineFunction &MF,
+                                 MachineFunctionAnalysisManager &) {
+  BlockUniformityProfile Profile;
+  Profile.compute(MF);
+  return Profile;
+}
+
+PreservedAnalyses
+BlockUniformityProfilePrinterPass::run(MachineFunction &MF,
+                                       MachineFunctionAnalysisManager &MFAM) {
+  auto &Profile = MFAM.getResult<BlockUniformityProfileProxy>(MF);
+  Profile.print(OS, MF);
+  return PreservedAnalyses::all();
+}
diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
index c572128b023c1..1305e12693fbd 100644
--- a/llvm/lib/CodeGen/CMakeLists.txt
+++ b/llvm/lib/CodeGen/CMakeLists.txt
@@ -36,6 +36,7 @@ add_llvm_component_library(LLVMCodeGen
   BasicBlockPathCloning.cpp
   BasicBlockSectionsProfileReader.cpp
   BasicBlockMatchingAndInference.cpp
+  BlockUniformityProfile.cpp
   CalcSpillWeights.cpp
   CallingConvLower.cpp
   CFGuardLongjmp.cpp
diff --git a/llvm/lib/CodeGen/SpillPlacement.cpp b/llvm/lib/CodeGen/SpillPlacement.cpp
index 55a96a22a00ec..4ea05ac2b2987 100644
--- a/llvm/lib/CodeGen/SpillPlacement.cpp
+++ b/llvm/lib/CodeGen/SpillPlacement.cpp
@@ -28,11 +28,14 @@
 
 #include "llvm/CodeGen/SpillPlacement.h"
 #include "llvm/ADT/BitVector.h"
+#include "llvm/CodeGen/BlockUniformityProfile.h"
 #include "llvm/CodeGen/EdgeBundles.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/Passes.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+#include "llvm/IR/Function.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include <algorithm>
@@ -193,7 +196,9 @@ bool SpillPlacementWrapperLegacy::runOnMachineFunction(MachineFunction &MF) {
   auto *Bundles = &getAnalysis<EdgeBundlesWrapperLegacy>().getEdgeBundles();
   auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
 
-  Impl.run(MF, Bundles, MBFI);
+  BlockUniformityProfile Profile;
+  Profile.compute(MF);
+  Impl.run(MF, Bundles, MBFI, &Profile);
   return false;
 }
 
@@ -204,8 +209,9 @@ SpillPlacementAnalysis::run(MachineFunction &MF,
                             MachineFunctionAnalysisManager &MFAM) {
   auto *Bundles = &MFAM.getResult<EdgeBundlesAnalysis>(MF);
   auto *MBFI = &MFAM.getResult<MachineBlockFrequencyAnalysis>(MF);
+  auto &Profile = MFAM.getResult<BlockUniformityProfileProxy>(MF);
   SpillPlacement Impl;
-  Impl.run(MF, Bundles, MBFI);
+  Impl.run(MF, Bundles, MBFI, &Profile);
   return Impl;
 }
 
@@ -217,7 +223,8 @@ bool SpillPlacementAnalysis::Result::invalidate(
     return true;
   // Check dependencies.
   return Inv.invalidate<EdgeBundlesAnalysis>(MF, PA) ||
-         Inv.invalidate<MachineBlockFrequencyAnalysis>(MF, PA);
+         Inv.invalidate<MachineBlockFrequencyAnalysis>(MF, PA) ||
+         Inv.invalidate<BlockUniformityProfileProxy>(MF, PA);
 }
 
 SpillPlacement::SpillPlacement() = default;
@@ -230,7 +237,8 @@ void SpillPlacement::releaseMemory() {
 }
 
 void SpillPlacement::run(MachineFunction &mf, EdgeBundles *Bundles,
-                         MachineBlockFrequencyInfo *MBFI) {
+                         MachineBlockFrequencyInfo *MBFI,
+                         const BlockUniformityProfile *Profile) {
   MF = &mf;
   this->bundles = Bundles;
   this->MBFI = MBFI;
@@ -240,12 +248,18 @@ void SpillPlacement::run(MachineFunction &mf, EdgeBundles *Bundles,
   TodoList.clear();
   TodoList.setUniverse(bundles->getNumBundles());
 
+  const bool HasProfile = Profile && Profile->hasProfile();
+
   // Compute total ingoing and outgoing block frequencies for all bundles.
   BlockFrequencies.resize(mf.getNumBlockIDs());
   setThreshold(MBFI->getEntryFreq());
   for (auto &I : mf) {
     unsigned Num = I.getNumber();
-    BlockFrequencies[Num] = MBFI->getBlockFreq(&I);
+    if (HasProfile && Profile->isDivergent(I)) {
+      BlockFrequencies[Num] = MBFI->getEntryFreq();
+    } else {
+      BlockFrequencies[Num] = MBFI->getBlockFreq(&I);
+    }
   }
 }
 
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 6dd9ab0eec054..6a2c5a09f2b92 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -82,6 +82,7 @@
 #include "llvm/CodeGen/AssignmentTrackingAnalysis.h"
 #include "llvm/CodeGen/AtomicExpand.h"
 #include "llvm/CodeGen/BasicBlockSectionsProfileReader.h"
+#include "llvm/CodeGen/BlockUniformityProfile.h"
 #include "llvm/CodeGen/BranchFoldingPass.h"
 #include "llvm/CodeGen/BranchRelaxation.h"
 #include "llvm/CodeGen/BreakFalseDeps.h"
diff --git a/llvm/test/CodeGen/AMDGPU/block-uniformity-profile.ll b/llvm/test/CodeGen/AMDGPU/block-uniformity-profile.ll
new file mode 100644
index 0000000000000..547f489143a51
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/block-uniformity-profile.ll
@@ -0,0 +1,79 @@
+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -O0 -stop-after=finalize-isel -o - %s | \
+; RUN:   llc -mtriple=amdgcn-amd-amdhsa -passes='print<block-uniformity-profile>' -x mir -filetype=null 2>&1 | FileCheck %s
+
+; Test that BlockUniformityProfileProxy correctly reads block.uniformity.profile
+; metadata from IR basic blocks and classifies machine blocks.
+;
+; This metadata is attached during PGO-use phase to indicate whether a basic block
+; was executed uniformly (all lanes together) or divergently (partial wave).
+;
+; The analysis is consumed by SpillPlacement to flatten block frequencies for
+; divergent blocks, preventing PGO from causing regressions on divergent code paths.
+
+; CHECK-LABEL: BlockUniformityProfile for function: @uniform_blocks
+; CHECK-NEXT: HasProfile: true
+; CHECK: %bb.{{[0-9]+}} (%entry): uniform
+define amdgpu_kernel void @uniform_blocks(ptr addrspace(1) %out) #0 {
+entry:
+  store i32 1, ptr addrspace(1) %out, align 4
+  ret void, !block.uniformity.profile !0
+}
+
+; CHECK-LABEL: BlockUniformityProfile for function: @divergent_blocks
+; CHECK-NEXT: HasProfile: true
+; CHECK-DAG: %bb.{{[0-9]+}} (%if.then): divergent
+; CHECK-DAG: %bb.{{[0-9]+}} (%if.else): uniform
+define amdgpu_kernel void @divergent_blocks(ptr addrspace(1) %out, i32 %tid) #0 {
+entry:
+  %cmp = icmp eq i32 %tid, 0
+  br i1 %cmp, label %if.then, label %if.else
+
+if.then:
+  store i32 1, ptr addrspace(1) %out, align 4
+  ret void, !block.uniformity.profile !1
+
+if.else:
+  store i32 2, ptr addrspace(1) %out, align 4
+  ret void, !block.uniformity.profile !0
+}
+
+; CHECK-LABEL: BlockUniformityProfile for function: @missing_metadata
+; CHECK-NEXT: HasProfile: true
+; CHECK-DAG: %bb.{{[0-9]+}} (%if.then): no PGO annotation (treated divergent for spill placement)
+; CHECK-DAG: %bb.{{[0-9]+}} (%if.else): uniform
+define amdgpu_kernel void @missing_metadata(ptr addrspace(1) %out, i32 %cond) #0 {
+entry:
+  %cmp = icmp sgt i32 %cond, 0
+  br i1 %cmp, label %if.then, label %if.else
+
+if.then:
+  store i32 1, ptr addrspace(1) %out, align 4
+  ret void
+
+if.else:
+  store i32 2, ptr addrspace(1) %out, align 4
+  ret void, !block.uniformity.profile !0
+}
+
+; CHECK-LABEL: BlockUniformityProfile for function: @no_divergence_metadata
+; CHECK-NEXT: HasProfile: false
+define amdgpu_kernel void @no_divergence_metadata(ptr addrspace(1) %out, i32 %cond) #0 {
+entry:
+  ; No uniformity metadata - analysis should report hasProfile() = false
+  %cmp = icmp sgt i32 %cond, 0
+  br i1 %cmp, label %if.then, label %if.else
+
+if.then:
+  store i32 1, ptr addrspace(1) %out, align 4
+  ret void
+
+if.else:
+  store i32 2, ptr addrspace(1) %out, align 4
+  ret void
+}
+
+attributes #0 = { "amdgpu-flat-work-group-size"="1,256" }
+
+; Metadata: i1 true = uniform, i1 false = divergent
+!0 = !{i1 true}   ; uniform
+!1 = !{i1 false}  ; divergent

>From 759e8cec8e8876e86d928254074c24cb3932fd6b Mon Sep 17 00:00:00 2001
From: "Yaxun (Sam) Liu" <yaxun.liu at amd.com>
Date: Thu, 2 Jul 2026 22:50:59 -0400
Subject: [PATCH 2/2] Address block uniformity profile review

---
 .../llvm/CodeGen/BlockUniformityProfile.h     |  2 +-
 llvm/lib/CodeGen/BlockUniformityProfile.cpp   | 47 +++++--------------
 .../block-uniformity-profile-metadata.ll      | 10 ++--
 .../AMDGPU/block-uniformity-profile.ll        | 35 +++++++++++---
 4 files changed, 46 insertions(+), 48 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/BlockUniformityProfile.h b/llvm/include/llvm/CodeGen/BlockUniformityProfile.h
index 209e0f893b6ca..d4802a0a5b092 100644
--- a/llvm/include/llvm/CodeGen/BlockUniformityProfile.h
+++ b/llvm/include/llvm/CodeGen/BlockUniformityProfile.h
@@ -10,7 +10,7 @@
 //
 // The source of truth is IR metadata attached during PGO use:
 //   - Metadata name: "block.uniformity.profile"
-//   - Payload: i1 (true = uniform, false = divergent)
+//   - Presence means the block is uniform.
 //
 // This is intentionally target-agnostic: any backend that produces
 // uniformity bits in the profile can attach the same metadata and reuse this
diff --git a/llvm/lib/CodeGen/BlockUniformityProfile.cpp b/llvm/lib/CodeGen/BlockUniformityProfile.cpp
index fff9843ce59eb..7a6400b9dee64 100644
--- a/llvm/lib/CodeGen/BlockUniformityProfile.cpp
+++ b/llvm/lib/CodeGen/BlockUniformityProfile.cpp
@@ -10,27 +10,16 @@
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/IR/BasicBlock.h"
-#include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/LLVMContext.h"
-#include "llvm/IR/Metadata.h"
 #include "llvm/Support/raw_ostream.h"
-#include <optional>
 
 using namespace llvm;
 
-static std::optional<bool> getIRBlockUniformity(const BasicBlock &BB) {
-  const Instruction *TI = BB.getTerminator();
-  if (!TI)
-    return std::nullopt;
-
-  MDNode *MD = TI->getMetadata(LLVMContext::MD_block_uniformity_profile);
-  if (!MD)
-    return std::nullopt;
-
-  // Metadata format: !{i1 IsUniform} - structural validity assumed (verifier).
-  // Returns true if uniform (not divergent).
-  return mdconst::extract<ConstantInt>(MD->getOperand(0))->isOne();
+static bool hasIRBlockUniformityProfile(const BasicBlock &BB) {
+  return BB.getTerminator()->getMetadata(
+      LLVMContext::MD_block_uniformity_profile);
 }
 
 void BlockUniformityProfile::compute(const MachineFunction &MF) {
@@ -39,30 +28,18 @@ void BlockUniformityProfile::compute(const MachineFunction &MF) {
   DivergentBlocks.clear();
   DivergentBlocks.resize(NumBlockIDs);
 
-  // First determine whether any uniformity profile exists for this function.
-  for (const MachineBasicBlock &MBB : MF) {
-    const BasicBlock *BB = MBB.getBasicBlock();
-    if (!BB)
-      continue;
-    if (getIRBlockUniformity(*BB).has_value()) {
-      HasProfile = true;
-      break;
-    }
-  }
-
-  if (!HasProfile)
-    return;
-
   // Conservative behavior: if profile exists for the function but we
   // cannot classify a particular (Machine)basic block, treat it as divergent.
   for (const MachineBasicBlock &MBB : MF) {
     const unsigned Num = MBB.getNumber();
-    bool IsDivergent = true;
+    bool IsUniform = false;
     if (const BasicBlock *BB = MBB.getBasicBlock()) {
-      if (auto U = getIRBlockUniformity(*BB))
-        IsDivergent = !*U; // Metadata stores IsUniform, we want IsDivergent
+      IsUniform = hasIRBlockUniformityProfile(*BB);
     }
-    if (Num < DivergentBlocks.size() && IsDivergent)
+    if (IsUniform)
+      HasProfile = true;
+
+    if (Num < DivergentBlocks.size() && !IsUniform)
       DivergentBlocks.set(Num);
   }
 }
@@ -83,8 +60,8 @@ void BlockUniformityProfile::print(raw_ostream &OS,
     OS << "  " << printMBBReference(MBB);
     if (BB->hasName())
       OS << " (%" << BB->getName() << ")";
-    if (auto U = getIRBlockUniformity(*BB)) {
-      OS << ": " << (*U ? "uniform" : "divergent") << '\n';
+    if (hasIRBlockUniformityProfile(*BB)) {
+      OS << ": uniform\n";
       continue;
     }
     OS << ": no PGO annotation (treated divergent for spill placement)\n";
diff --git a/llvm/test/Bitcode/block-uniformity-profile-metadata.ll b/llvm/test/Bitcode/block-uniformity-profile-metadata.ll
index b69956fa26622..92ca4c7c72680 100644
--- a/llvm/test/Bitcode/block-uniformity-profile-metadata.ll
+++ b/llvm/test/Bitcode/block-uniformity-profile-metadata.ll
@@ -11,11 +11,9 @@ uniform:
   ret void
 
 divergent:
-  br label %uniform, !block.uniformity.profile !1
-; CHECK: br label %uniform, !block.uniformity.profile !1
+  br label %uniform, !block.uniformity.profile !0
+; CHECK: br label %uniform, !block.uniformity.profile !0
 }
 
-; CHECK: !0 = !{i1 true}
-; CHECK: !1 = !{i1 false}
-!0 = !{i1 true}
-!1 = !{i1 false}
+; CHECK: !0 = !{}
+!0 = !{}
diff --git a/llvm/test/CodeGen/AMDGPU/block-uniformity-profile.ll b/llvm/test/CodeGen/AMDGPU/block-uniformity-profile.ll
index 547f489143a51..93297aca5886d 100644
--- a/llvm/test/CodeGen/AMDGPU/block-uniformity-profile.ll
+++ b/llvm/test/CodeGen/AMDGPU/block-uniformity-profile.ll
@@ -5,7 +5,8 @@
 ; metadata from IR basic blocks and classifies machine blocks.
 ;
 ; This metadata is attached during PGO-use phase to indicate whether a basic block
-; was executed uniformly (all lanes together) or divergently (partial wave).
+; was executed uniformly (all lanes together). Missing metadata in a profiled
+; function is conservatively treated as divergent.
 ;
 ; The analysis is consumed by SpillPlacement to flatten block frequencies for
 ; divergent blocks, preventing PGO from causing regressions on divergent code paths.
@@ -21,7 +22,7 @@ entry:
 
 ; CHECK-LABEL: BlockUniformityProfile for function: @divergent_blocks
 ; CHECK-NEXT: HasProfile: true
-; CHECK-DAG: %bb.{{[0-9]+}} (%if.then): divergent
+; CHECK-DAG: %bb.{{[0-9]+}} (%if.then): no PGO annotation (treated divergent for spill placement)
 ; CHECK-DAG: %bb.{{[0-9]+}} (%if.else): uniform
 define amdgpu_kernel void @divergent_blocks(ptr addrspace(1) %out, i32 %tid) #0 {
 entry:
@@ -30,7 +31,7 @@ entry:
 
 if.then:
   store i32 1, ptr addrspace(1) %out, align 4
-  ret void, !block.uniformity.profile !1
+  ret void
 
 if.else:
   store i32 2, ptr addrspace(1) %out, align 4
@@ -55,6 +56,29 @@ if.else:
   ret void, !block.uniformity.profile !0
 }
 
+; CHECK-LABEL: BlockUniformityProfile for function: @loop_blocks
+; CHECK-NEXT: HasProfile: true
+; CHECK-DAG: %bb.{{[0-9]+}} (%loop.header): no PGO annotation (treated divergent for spill placement)
+; CHECK-DAG: %bb.{{[0-9]+}} (%loop.body): no PGO annotation (treated divergent for spill placement)
+; CHECK-DAG: %bb.{{[0-9]+}} (%exit): uniform
+define amdgpu_kernel void @loop_blocks(ptr addrspace(1) %out, i32 %n) #0 {
+entry:
+  br label %loop.header
+
+loop.header:
+  %i = phi i32 [ 0, %entry ], [ %inc, %loop.body ]
+  %cmp = icmp slt i32 %i, %n
+  br i1 %cmp, label %loop.body, label %exit
+
+loop.body:
+  store i32 %i, ptr addrspace(1) %out, align 4
+  %inc = add nuw nsw i32 %i, 1
+  br label %loop.header
+
+exit:
+  ret void, !block.uniformity.profile !0
+}
+
 ; CHECK-LABEL: BlockUniformityProfile for function: @no_divergence_metadata
 ; CHECK-NEXT: HasProfile: false
 define amdgpu_kernel void @no_divergence_metadata(ptr addrspace(1) %out, i32 %cond) #0 {
@@ -74,6 +98,5 @@ if.else:
 
 attributes #0 = { "amdgpu-flat-work-group-size"="1,256" }
 
-; Metadata: i1 true = uniform, i1 false = divergent
-!0 = !{i1 true}   ; uniform
-!1 = !{i1 false}  ; divergent
+; Metadata presence means uniform.
+!0 = !{}



More information about the llvm-commits mailing list