From 3df59d9ae7282afb0e4326441c9dc4b7de82968d Mon Sep 17 00:00:00 2001 From: Mircea Trofin Date: Mon, 7 Jul 2025 12:44:03 -0700 Subject: [PATCH] pass validator --- .../llvm/Transforms/Utils/ProfileVerify.h | 36 ++++++ llvm/lib/Passes/PassBuilder.cpp | 1 + llvm/lib/Passes/PassRegistry.def | 2 + llvm/lib/Transforms/Utils/CMakeLists.txt | 1 + llvm/lib/Transforms/Utils/ProfileVerify.cpp | 117 ++++++++++++++++++ .../PGOProfile/prof-verify-as-needed.ll | 20 +++ .../PGOProfile/prof-verify-existing.ll | 21 ++++ .../test/Transforms/PGOProfile/prof-verify.ll | 19 +++ 8 files changed, 217 insertions(+) create mode 100644 llvm/include/llvm/Transforms/Utils/ProfileVerify.h create mode 100644 llvm/lib/Transforms/Utils/ProfileVerify.cpp create mode 100644 llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll create mode 100644 llvm/test/Transforms/PGOProfile/prof-verify-existing.ll create mode 100644 llvm/test/Transforms/PGOProfile/prof-verify.ll diff --git a/llvm/include/llvm/Transforms/Utils/ProfileVerify.h b/llvm/include/llvm/Transforms/Utils/ProfileVerify.h new file mode 100644 index 0000000000000..7834305b569d2 --- /dev/null +++ b/llvm/include/llvm/Transforms/Utils/ProfileVerify.h @@ -0,0 +1,36 @@ +//===- ProfileVerify.h - Verify profile info for testing ----------*-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 +// +//===----------------------------------------------------------------------===// +// +// Inject profile information, as part of tests, to verify passes don't +// accidentally drop it. +// +//===----------------------------------------------------------------------===// +#ifndef LLVM_TRANSFORMS_UTILS_PROFILEVERIFY_H +#define LLVM_TRANSFORMS_UTILS_PROFILEVERIFY_H + +#include "llvm/IR/Analysis.h" +#include "llvm/IR/PassManager.h" + +namespace llvm { +/// Inject MD_prof metadata where it's missing. Used for testing that passes +/// don't accidentally drop this metadata. +class ProfileInjectorPass : public PassInfoMixin { +public: + PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM); +}; + +/// Checks that MD_prof is present on every instruction that supports it. Used +/// in conjunction with the ProfileInjectorPass. MD_prof "unknown" is considered +/// valid (i.e. !{!"unknown"}) +class ProfileVerifierPass : public PassInfoMixin { +public: + PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM); +}; + +} // namespace llvm +#endif diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp index 874fce05841e2..70c0e999cbb95 100644 --- a/llvm/lib/Passes/PassBuilder.cpp +++ b/llvm/lib/Passes/PassBuilder.cpp @@ -357,6 +357,7 @@ #include "llvm/Transforms/Utils/MoveAutoInit.h" #include "llvm/Transforms/Utils/NameAnonGlobals.h" #include "llvm/Transforms/Utils/PredicateInfo.h" +#include "llvm/Transforms/Utils/ProfileVerify.h" #include "llvm/Transforms/Utils/RelLookupTableConverter.h" #include "llvm/Transforms/Utils/StripGCRelocates.h" #include "llvm/Transforms/Utils/StripNonLineTableDebugInfo.h" diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def index dd3dab3425975..dfe233f15a149 100644 --- a/llvm/lib/Passes/PassRegistry.def +++ b/llvm/lib/Passes/PassRegistry.def @@ -517,6 +517,8 @@ FUNCTION_PASS("print", RegionInfoPrinterPass(errs())) FUNCTION_PASS("print", ScalarEvolutionPrinterPass(errs())) FUNCTION_PASS("print", StackSafetyPrinterPass(errs())) FUNCTION_PASS("print", UniformityInfoPrinterPass(errs())) +FUNCTION_PASS("prof-inject", ProfileInjectorPass()) +FUNCTION_PASS("prof-verify", ProfileVerifierPass()) FUNCTION_PASS("reassociate", ReassociatePass()) FUNCTION_PASS("redundant-dbg-inst-elim", RedundantDbgInstEliminationPass()) FUNCTION_PASS("reg2mem", RegToMemPass()) diff --git a/llvm/lib/Transforms/Utils/CMakeLists.txt b/llvm/lib/Transforms/Utils/CMakeLists.txt index 78cad0d253be8..c0bd6d647aad4 100644 --- a/llvm/lib/Transforms/Utils/CMakeLists.txt +++ b/llvm/lib/Transforms/Utils/CMakeLists.txt @@ -67,6 +67,7 @@ add_llvm_component_library(LLVMTransformUtils MoveAutoInit.cpp NameAnonGlobals.cpp PredicateInfo.cpp + ProfileVerify.cpp PromoteMemoryToRegister.cpp RelLookupTableConverter.cpp ScalarEvolutionExpander.cpp diff --git a/llvm/lib/Transforms/Utils/ProfileVerify.cpp b/llvm/lib/Transforms/Utils/ProfileVerify.cpp new file mode 100644 index 0000000000000..54cf1397b5f9a --- /dev/null +++ b/llvm/lib/Transforms/Utils/ProfileVerify.cpp @@ -0,0 +1,117 @@ +//===- ProfileVerify.cpp - Verify profile info for testing ----------------===// +// +// 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/Transforms/Utils/ProfileVerify.h" +#include "llvm/ADT/DynamicAPInt.h" +#include "llvm/ADT/PostOrderIterator.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/Analysis/BranchProbabilityInfo.h" +#include "llvm/Analysis/LoopInfo.h" +#include "llvm/IR/Analysis.h" +#include "llvm/IR/Dominators.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/MDBuilder.h" +#include "llvm/IR/ProfDataUtils.h" +#include "llvm/Support/BranchProbability.h" + +using namespace llvm; +namespace { +class ProfileInjector { + Function &F; + FunctionAnalysisManager &FAM; + +public: + static const Instruction * + getTerminatorBenefitingFromMDProf(const BasicBlock &BB) { + if (succ_size(&BB) < 2) + return nullptr; + auto *Term = BB.getTerminator(); + return (isa(Term) || isa(Term) || + isa(Term) || isa(Term)) + ? Term + : nullptr; + } + + static Instruction *getTerminatorBenefitingFromMDProf(BasicBlock &BB) { + return const_cast( + getTerminatorBenefitingFromMDProf(const_cast(BB))); + } + + ProfileInjector(Function &F, FunctionAnalysisManager &FAM) : F(F), FAM(FAM) {} + bool inject(); +}; +} // namespace + +// FIXME: currently this injects only for terminators. Select isn't yet +// supported. +bool ProfileInjector::inject() { + auto &BPI = FAM.getResult(F); + + bool Changed = false; + for (auto &BB : F) { + auto *Term = getTerminatorBenefitingFromMDProf(BB); + if (!Term || Term->getMetadata(LLVMContext::MD_prof)) + continue; + SmallVector Probs; + Probs.reserve(Term->getNumSuccessors()); + for (auto I = 0U, E = Term->getNumSuccessors(); I < E; ++I) + Probs.emplace_back(BPI.getEdgeProbability(&BB, Term->getSuccessor(I))); + + const auto *FirstZeroDenominator = + find_if(Probs, [](const BranchProbability &P) { + return P.getDenominator() == 0; + }); + (void)FirstZeroDenominator; + assert(FirstZeroDenominator == Probs.end()); + const auto *FirstNonZeroNumerator = + find_if(Probs, [](const BranchProbability &P) { + return !P.isZero() && !P.isUnknown(); + }); + assert(FirstNonZeroNumerator != Probs.end()); + DynamicAPInt LCM(Probs[0].getDenominator()); + DynamicAPInt GCD(FirstNonZeroNumerator->getNumerator()); + for (const auto &Prob : drop_begin(Probs)) { + if (!Prob.getNumerator()) + continue; + LCM = llvm::lcm(LCM, DynamicAPInt(Prob.getDenominator())); + GCD = llvm::gcd(GCD, DynamicAPInt(Prob.getNumerator())); + } + SmallVector Weights; + Weights.reserve(Term->getNumSuccessors()); + for (const auto &Prob : Probs) { + DynamicAPInt W = + (Prob.getNumerator() * LCM / GCD) / Prob.getDenominator(); + Weights.emplace_back(static_cast((int64_t)W)); + } + setBranchWeights(*Term, Weights, /*IsExpected=*/false); + Changed = true; + } + return Changed; +} + +PreservedAnalyses ProfileInjectorPass::run(Function &F, + FunctionAnalysisManager &FAM) { + ProfileInjector PI(F, FAM); + if (!PI.inject()) + return PreservedAnalyses::all(); + + return PreservedAnalyses::none(); +} + +PreservedAnalyses ProfileVerifierPass::run(Function &F, + FunctionAnalysisManager &FAM) { + for (const auto &BB : F) + if (const auto *Term = + ProfileInjector::getTerminatorBenefitingFromMDProf(BB)) + if (!Term->getMetadata(LLVMContext::MD_prof)) + F.getContext().emitError("Profile verification failed"); + + return PreservedAnalyses::none(); +} diff --git a/llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll b/llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll new file mode 100644 index 0000000000000..07e1f2d3c6127 --- /dev/null +++ b/llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll @@ -0,0 +1,20 @@ +; Test that prof-inject only injects missing metadata + +; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s + +define void @foo(i32 %i) { + %c = icmp eq i32 %i, 0 + br i1 %c, label %yes, label %no, !prof !0 +yes: + br i1 %c, label %yes2, label %no +yes2: + ret void +no: + ret void +} + +!0 = !{!"branch_weights", i32 1, i32 2} +; CHECK: br i1 %c, label %yes, label %no, !prof !0 +; CHECK: br i1 %c, label %yes2, label %no, !prof !1 +; CHECK: !0 = !{!"branch_weights", i32 1, i32 2} +; CHECK: !1 = !{!"branch_weights", i32 3, i32 5} diff --git a/llvm/test/Transforms/PGOProfile/prof-verify-existing.ll b/llvm/test/Transforms/PGOProfile/prof-verify-existing.ll new file mode 100644 index 0000000000000..ea4f0f9f1dadf --- /dev/null +++ b/llvm/test/Transforms/PGOProfile/prof-verify-existing.ll @@ -0,0 +1,21 @@ +; Test that prof-inject does not modify existing metadata (incl. "unknown") + +; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s +; RUN: opt -passes=prof-verify %s -S --disable-output + +define void @foo(i32 %i) { + %c = icmp eq i32 %i, 0 + br i1 %c, label %yes, label %no, !prof !0 +yes: + br i1 %c, label %yes2, label %no, !prof !1 +yes2: + ret void +no: + ret void +} + +!0 = !{!"branch_weights", i32 1, i32 2} +!1 = !{!"unknown"} +; CHECK: br i1 %c, label %yes, label %no, !prof !0 +; CHECK: !0 = !{!"branch_weights", i32 1, i32 2} +; CHECK: !1 = !{!"unknown"} diff --git a/llvm/test/Transforms/PGOProfile/prof-verify.ll b/llvm/test/Transforms/PGOProfile/prof-verify.ll new file mode 100644 index 0000000000000..c83475ab5d18e --- /dev/null +++ b/llvm/test/Transforms/PGOProfile/prof-verify.ll @@ -0,0 +1,19 @@ +; Test prof-inject and prof-verify + +; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s --check-prefix=INJECT +; RUN: not opt -passes=prof-verify %s -S -o - 2>&1 | FileCheck %s --check-prefix=VERIFY +; RUN: opt -passes=prof-inject,prof-verify %s --disable-output + +define void @foo(i32 %i) { + %c = icmp eq i32 %i, 0 + br i1 %c, label %yes, label %no +yes: + ret void +no: + ret void +} + +; INJECT: br i1 %c, label %yes, label %no, !prof !0 +; INJECT: !0 = !{!"branch_weights", i32 3, i32 5} + +; VERIFY: Profile verification failed \ No newline at end of file