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

[Yaxun Liu via llvm-commits]

================
@@ -0,0 +1,117 @@
+//===- OffloadBlockUniformity.cpp - Offload block uniformity info --------===//
+//
+// 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/OffloadBlockUniformity.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/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(OffloadBlockUniformityInfo::MetadataName);
+  if (!MD)
+    return std::nullopt;
+
+  // Metadata format: !{i1 IsUniform} - structural validity assumed (verifier).
+  return mdconst::extract<ConstantInt>(MD->getOperand(0))->isOne();
+}
+
+void OffloadBlockUniformityInfo::compute(const MachineFunction &MF) {
+  HasAnyUniformity = false;
+  DivergentBlocks.clear();
+  DivergentBlocks.resize(MF.getNumBlockIDs());
+
+  // First determine whether any uniformity annotation exists for this function.
+  for (const MachineBasicBlock &MBB : MF) {
+    const BasicBlock *BB = MBB.getBasicBlock();
+    if (!BB)
+      continue;
+    if (getIRBlockUniformity(*BB).has_value()) {
+      HasAnyUniformity = true;
+      break;
+    }
+  }
+
+  if (!HasAnyUniformity)
+    return;
+
+  // Conservative behavior: if uniformity 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;
+    }
+    if (Num < DivergentBlocks.size() && IsDivergent)
----------------
yxsamliu wrote:

good catch. I think normally this should not happen unless some pass changed MBB numbers but does not invalidate the analysis. will add assert that the size of the bit vector matches number of MBB's.

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