[llvm] Provide utility function for MD_prof

[paulkirth]

Summary

Currently, there is significant code duplication for dealing with
MD_prof metadata throughout the compiler. These utility functions can
improve code reuse and simplify boilerplate code when dealing with
profiling metadata, such as branch weights. The inent is to provide a
uniform set of APIs that allow common tasks, such as identifying
specific types of MD_prof metadata and extracting branch weights.

Future patches can build on this initial implementation and clean up the
different implementations across the compiler.

[paulkirth]

Simplify helper functions and remove helpers for MD_prof types other than branch_weights

[paulkirth]

Refactor interfaces

[paulkirth]

Add assert back

[paulkirth]

Small fixup from rebase

[paulkirth]

There's an argument to be made that most of these utility functions should just be part of Metadata.cpp. Separating them seems like a good idea, but as I've implemented this and begun to integrate their use in D128860, I've found that there seems to be significant overlap w/ some of the utility functions in Metadata.cpp

I'll defer to reviewers on precisely where this should live, since there are a few good alternatives in my mind.

[bogner]

There's an argument to be made that most of these utility functions should just be part of Metadata.cpp.

It's easy enough to merge these into Metadata.cpp later if we want to. I think a separate utils library is fine for now.

llvm/include/llvm/IR/ProfDataUtils.h

43

typo totla/total

47

Would "extractRawBranchWeights" be a better name here? I guess this matches the naming of the user, but I think it's a bit hard to tell the difference between this and the extractBranchWeights functions from the name/doc comments as is.

llvm/lib/IR/ProfDataUtils.cpp

51

Can we actually get here if we have valid metadata? This looks like something is malformed

[paulkirth]

llvm/include/llvm/IR/ProfDataUtils.h

47

Maybe we should leave this one as is, and rename the others to extractRawBranchWeights? This one seems really clear to me that its extracting the total of all weights.

One wrinkle here is that to me extractRawBranchWeights communicates more that we're extracting a series of weights in their raw form. Under that reading I'd expect it to be a SmallVector<ConstantInt*>, i.e., the raw format in the metadata.

In the meantime I'm going to look around the codebase a bit and see if I can either find a pattern we can use or think of another more communicative name.

llvm/lib/IR/ProfDataUtils.cpp

51

That may be true but most of the places where we're extracting branch weights have similar checks.

Two examples are:
https://github.com/llvm/llvm-project/blob/3a2d7d8ad5291102ad98b8269b9c2171d3d58d76/llvm/lib/Analysis/BranchProbabilityInfo.cpp#L404
https://github.com/llvm/llvm-project/blob/3a2d7d8ad5291102ad98b8269b9c2171d3d58d76/llvm/lib/IR/Metadata.cpp#L1510

It looks like we didn't follow that pattern in Transform/Utils/MisExpect.cpp, but that may be a mistake on my part.

I can seen an argument in both cases, and I'm not sure one is better than the other. Since I don't have a strong opinion, I'm happy to replace the check with an assert though if you think that's the better choice.

[paulkirth]

respond to comments

• Modify API name for extracting the total weight to match the existing API.
• Add specialized version of extractBranchWeights for branches and selects to ease replacing uses in LLVM.
• Replace checks on values of metadata w/ asserts that the metadata is not malformed

[paulkirth]

cleanup and formatting

llvm/include/llvm/IR/ProfDataUtils.h

+#ifndef PROF_DATA_UTILS_H
+#define PROF_DATA_UTILS_H
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/IR/Metadata.h"
+
+namespace llvm {
+
+/// Checks if an Instruction has MD_prof Metadata
+bool hasProfMD(const Instruction &I);
+
+/// Checks if an MDNode contains Branch Weight Metadata
+bool isBranchWeightMD(const MDNode *ProfileData);
+
+/// Checks if an instructions has Branch Weight Metadata
+///
+/// \param I The instruction to check
+/// \return True if I has an MD_prof node containing Branch Weights. False
+/// otherwise.
+bool hasBranchWeightMD(const Instruction &I);
+
+/// Extract branch weights from MD_prof metadata
+///
+/// \param ProfileData A pointer to an MDNode.
+/// \param Weights An output vector to fill with branch weights
+/// \return True if weights were extracted, False otherwise. When false Weights
+/// will be cleared.
+bool extractBranchWeights(const MDNode *ProfileData,
+                          SmallVectorImpl<uint32_t> &Weights);
+
+/// Extract branch weights attatched to an Instruction
+///
+/// \param I The Instruction to extract weights from.
+/// \param Weights An output vector to fill with branch weights
+/// \return True if weights were extracted, False otherwise. When false Weights
+/// will be cleared.
+bool extractBranchWeights(const Instruction &I,
+                          SmallVectorImpl<uint32_t> &Weights);
+
+/// Retrieve the raw weight values of a conditional branch or select.
+/// Returns true on success with profile weights filled in.
+/// Returns false if no metadata or invalid metadata was found.
+bool extractBranchWeights(const Instruction &I, uint64_t &TrueVal,
+                          uint64_t &FalseVal);
+
+/// Retrieve the total of all weights from MD_prof data.
+///
+/// \param ProfileData The profile data to extract the total weight from
+/// \param TotalWeights input variable to fill with total weights
+/// \return true on success with profile total weights filled in.
+/// \return false if no metadata was found.
+bool extractProfTotalWeight(const MDNode *ProfileData, uint64_t &TotalWeights);
+
+} // namespace llvm
+#endif

llvm/lib/IR/CMakeLists.txt

add_llvm_component_library(LLVMCore
   Operator.cpp
   OptBisect.cpp
   Pass.cpp
   PassInstrumentation.cpp
   PassManager.cpp
   PassRegistry.cpp
   PassTimingInfo.cpp
   PrintPasses.cpp
+  ProfDataUtils.cpp
   SafepointIRVerifier.cpp
   ProfileSummary.cpp
   PseudoProbe.cpp
   ReplaceConstant.cpp
   Statepoint.cpp
   StructuralHash.cpp
   Type.cpp
   TypeFinder.cpp

llvm/lib/IR/ProfDataUtils.cpp

+#include "llvm/IR/ProfDataUtils.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/Support/BranchProbability.h"
+#include "llvm/Support/CommandLine.h"
+
+using namespace llvm;
+
+namespace {
+
+// MD_prof nodes have the following layout
+//
+// In general:
+// { String name,         Array of i32   }
+//
+// In terms of Types:
+// { MDString,            [i32, i32, ...]}
+//
+// Concretely for Branch Weights
+// { "branch_weights",    [i32 1, i32 10000]}
+//
+// We maintain some constants here to ensure that we access the branch weights
+// correctly, and can change the behavior in the future if the layout changes
+
+// The index at which the weights vector starts
+constexpr unsigned WeightsIdx = 1;
+
+// the minimum number of operands for MD_prof nodes with branch weights
+constexpr unsigned MinBWOps = 3;
+
+bool extractWeights(const MDNode *ProfileData,
+                    SmallVectorImpl<uint32_t> &Weights) {
+  // Assume preconditions are already met (i.e. this is valid metadata)
+  assert(ProfileData && "ProfileData was nullptr in extractWeights");
+  unsigned NOps = ProfileData->getNumOperands();
+
+  assert(WeightsIdx < NOps && "Weights Index must be less than NOps.");
+  Weights.resize(NOps - WeightsIdx);
+
+  for (unsigned Idx = WeightsIdx, E = NOps; Idx != E; ++Idx) {
+    ConstantInt *Weight =
+        mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(Idx));
+    assert(Weight && "Malformed branch_weight in MD_prof node");
+    assert(Weight->getValue().getActiveBits() <= 32 &&
+           "Too many bits for uint32_t");
+    Weights[Idx - WeightsIdx] = Weight->getZExtValue();
+  }
+  return true;
+}
+
+// We may want to add support for other MD_prof types, so provide an abstraction
+// for checking the metadata type.
+bool isTargetMD(const MDNode *ProfData, const char *Name, unsigned MinOps) {
+  // TODO: This routine may be simplified if MD_prof used an enum instead of a
+  // string to differentiate the types of MD_prof nodes.
+  if (!ProfData || !Name || MinOps < 2)
+    return false;
+
+  unsigned NOps = ProfData->getNumOperands();
+  if (NOps < MinOps)
+    return false;
+
+  auto *ProfDataName = dyn_cast<MDString>(ProfData->getOperand(0));
+  if (!ProfDataName)
+    return false;
+
+  return ProfDataName->getString().equals(Name);
+}
+
+} // namespace
+
+namespace llvm {
+
+bool hasProfMD(const Instruction &I) {
+  return nullptr != I.getMetadata(LLVMContext::MD_prof);
+}
+
+bool isBranchWeightMD(const MDNode *ProfileData) {
+  return isTargetMD(ProfileData, "branch_weights", MinBWOps);
+}
+
+bool hasBranchWeightMD(const Instruction &I) {
+  auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
+  return isBranchWeightMD(ProfileData);
+}
+
+bool extractBranchWeights(const MDNode *ProfileData,
+                          SmallVectorImpl<uint32_t> &Weights) {
+  if (!isBranchWeightMD(ProfileData))
+    return false;
+  return extractWeights(ProfileData, Weights);
+}
+
+bool extractBranchWeights(const Instruction &I,
+                          SmallVectorImpl<uint32_t> &Weights) {
+  auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
+  return extractBranchWeights(ProfileData, Weights);
+}
+
+bool extractBranchWeights(const Instruction &I, uint64_t &TrueVal,
+                          uint64_t &FalseVal) {
+  assert((I.getOpcode() == Instruction::Br ||
+          I.getOpcode() == Instruction::Select) &&
+         "Looking for branch weights on something besides branch or select");
+
+  SmallVector<uint32_t, 2> Weights;
+  auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
+  if (!extractBranchWeights(ProfileData, Weights))
+    return false;
+
+  if (Weights.size() > 2)
+    return false;
+
+  TrueVal = Weights[0];
+  FalseVal = Weights[1];
+  return true;
+}
+
+bool extractProfTotalWeight(const MDNode *ProfileData, uint64_t &TotalVal) {
+  TotalVal = 0;
+  if (!ProfileData)
+    return false;
+
+  auto *ProfDataName = dyn_cast<MDString>(ProfileData->getOperand(0));
+  if (!ProfDataName)
+    return false;
+
+  if (ProfDataName->getString().equals("branch_weights")) {
+    for (unsigned Idx = 1; Idx < ProfileData->getNumOperands(); Idx++) {
+      auto *V = mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(Idx));
+      assert(V && "Malformed branch_weight in MD_prof node");
+      TotalVal += V->getValue().getZExtValue();
+    }
+    return true;
+  } else if (ProfDataName->getString().equals("VP") &&
+             ProfileData->getNumOperands() > 3) {
+    TotalVal = mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(2))
+                   ->getValue()
+                   .getZExtValue();
+    return true;
+  }
+  return false;
+}
+
+} // namespace llvm

llvm/unittests/Transforms/Utils/CMakeLists.txt

add_llvm_unittest(UtilsTests
   MemTransferLowering.cpp
   ModuleUtilsTest.cpp
   ScalarEvolutionExpanderTest.cpp
   SizeOptsTest.cpp
   SSAUpdaterBulkTest.cpp
   UnrollLoopTest.cpp
   ValueMapperTest.cpp
   VFABIUtils.cpp
+  ProfDataUtilTest.cpp
   )
 
 set_property(TARGET UtilsTests PROPERTY FOLDER "Tests/UnitTests/TransformsTests")

llvm/unittests/Transforms/Utils/ProfDataUtilTest.cpp

+//===----- ProfDataUtils.cpp - Unit tests for ProfDataUtils ---------------===//
+//
+// 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/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/BasicAliasAnalysis.h"
+#include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
+#include "llvm/Analysis/CFG.h"
+#include "llvm/Analysis/DomTreeUpdater.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/MemorySSA.h"
+#include "llvm/Analysis/MemorySSAUpdater.h"
+#include "llvm/Analysis/PostDominators.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/ProfDataUtils.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Transforms/Utils/BreakCriticalEdges.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+static std::unique_ptr<Module> parseIR(LLVMContext &C, const char *IR) {
+  SMDiagnostic Err;
+  std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);
+  if (!Mod)
+    Err.print("ProfDataUtilsTests", errs());
+  return Mod;
+}
+
+TEST(ProfDataUtils, extractWeights) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = parseIR(C, R"IR(
+define void @foo(i1 %cond0) {
+entry:
+  br i1 %cond0, label %bb0, label %bb1, !prof !1
+bb0:
+ %0 = mul i32 1, 2
+ br label %bb1
+bb1:
+  ret void
+}
+
+!1 = !{!"branch_weights", i32 1, i32 100000}
+)IR");
+  Function *F = M->getFunction("foo");
+  auto &Entry = F->getEntryBlock();
+  auto &I = Entry.front();
+  auto *Branch = dyn_cast<BranchInst>(&I);
+  EXPECT_NE(nullptr, Branch);
+  auto *ProfileData = Branch->getMetadata(LLVMContext::MD_prof);
+  EXPECT_NE(ProfileData, nullptr);
+  EXPECT_TRUE(hasProfMD(I));
+  SmallVector<uint32_t> Weights;
+  EXPECT_TRUE(extractBranchWeights(ProfileData, Weights));
+  EXPECT_EQ(Weights[0], 1U);
+  EXPECT_EQ(Weights[1], 100000U);
+  EXPECT_EQ(Weights.size(), 2U);
+}
+
+TEST(ProfDataUtils, NoWeights) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = parseIR(C, R"IR(
+define void @foo(i1 %cond0) {
+entry:
+  br i1 %cond0, label %bb0, label %bb1
+bb0:
+ %0 = mul i32 1, 2
+ br label %bb1
+bb1:
+  ret void
+}
+)IR");
+  Function *F = M->getFunction("foo");
+  auto &Entry = F->getEntryBlock();
+  auto &I = Entry.front();
+  auto *Branch = dyn_cast<BranchInst>(&I);
+  EXPECT_NE(nullptr, Branch);
+  auto *ProfileData = Branch->getMetadata(LLVMContext::MD_prof);
+  EXPECT_EQ(ProfileData, nullptr);
+  EXPECT_FALSE(hasProfMD(I));
+  SmallVector<uint32_t> Weights;
+  EXPECT_FALSE(extractBranchWeights(ProfileData, Weights));
+  EXPECT_EQ(Weights.size(), 0U);
+}