Plugin.cpp

#include "llvm/Passes/PassBuilder.h"
#include "llvm/Passes/PassPlugin.h"

using namespace llvm;

struct MyAnalysis : AnalysisInfoMixin<MyAnalysis> {
    // The result type the analysis produces
    struct Result {
        size_t numStatements = 0;

        bool invalidate(Module& IR, const PreservedAnalyses& PA, ModuleAnalysisManager::Invalidator& Inv) {
            return false;
        }

        bool invalidate(Function& IR, const PreservedAnalyses& PA, FunctionAnalysisManager::Invalidator& Inv) {
            return false;
        }
    };

    // Module entry point
    Result run(Module& M, ModuleAnalysisManager&) {
        Result R;
        for (Function& F : M) {
            if (F.isDeclaration()) { continue; }
            R.numStatements += F.getInstructionCount();
        }
        return R;
    }

    // Function entry point
    Result run(Function& F, FunctionAnalysisManager&) {
        Result R;
        if (F.isDeclaration()) return R;
        R.numStatements += F.getInstructionCount();
        return R;
    }

private:
    // Required for AnalysisInfoMixin
    friend AnalysisInfoMixin<MyAnalysis>;
    static AnalysisKey Key;
};

AnalysisKey MyAnalysis::Key;


struct IRDemoPlugin : PassInfoMixin<IRDemoPlugin> {
    PreservedAnalyses run(Function& function, FunctionAnalysisManager& FA) {
        outs() << "Running IRDemoPass on function " << function.getName() << "\n";
        MyAnalysis::Result functionLevelResult = FA.getResult<MyAnalysis>(function);
        outs() << "Function has " << functionLevelResult.numStatements << " instructions\n";
        MyAnalysis::Result* moduleLevelResult = FA.getResult<ModuleAnalysisManagerFunctionProxy>(function).
                                                   getCachedResult<MyAnalysis>(*function.getParent());
        if (!moduleLevelResult) {
            outs() << "ModuleLevelResult is unavailable\n";
            return PreservedAnalyses::all();;
        }

        outs() << "This is " << format(
                "%.2f", (functionLevelResult.numStatements / (double)moduleLevelResult->numStatements * 100)) <<
            "% of all instructions\n";

        return PreservedAnalyses::all();
    }
};

PassPluginLibraryInfo getPluginInfo() {
    return {
        LLVM_PLUGIN_API_VERSION, "DemoPlugin", "1.0", [](PassBuilder& PB) {
            // allow registration via optlevel (non-lto and thin-lto)
            PB.registerOptimizerEarlyEPCallback(
                [](ModulePassManager& MPM, OptimizationLevel, ThinOrFullLTOPhase phase) {
                    switch (phase) {
                    case ThinOrFullLTOPhase::None: outs() << "Registering IRDemoPlugin to run during opt\n";
                        break;
                    case ThinOrFullLTOPhase::ThinLTOPreLink: outs() <<
                            "Registering IRDemoPlugin to run during ThinLTOPreLink\n";
                        break;
                    case ThinOrFullLTOPhase::ThinLTOPostLink: outs() <<
                            "Registering IRDemoPlugin to run during ThinLTOPostLink\n";
                        break;
                    case ThinOrFullLTOPhase::FullLTOPreLink: outs() <<
                            "Registering IRDemoPlugin to run during FullLTOPreLink\n";
                        break;
                    case ThinOrFullLTOPhase::FullLTOPostLink: outs() <<
                            "Registering IRDemoPlugin to run during FullLTOPostLink\n";
                        break;
                    default: __builtin_unreachable();
                    }

                    FunctionPassManager FPM;
                    FPM.addPass(IRDemoPlugin());
                    MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
                });

            // allow registration via optlevel during full-lto
            PB.registerFullLinkTimeOptimizationLastEPCallback([](ModulePassManager& MPM, OptimizationLevel o) {
                outs() << "Registering IRDemoPlugin to run during full LTO\n";
                FunctionPassManager FPM;
                FPM.addPass(IRDemoPlugin());
                // Wrap function pass into module pass
                MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
            });

            // allow registration via opt pipeline parser
            PB.registerPipelineParsingCallback(
                [](StringRef Name, FunctionPassManager& FPM, ArrayRef<PassBuilder::PipelineElement>) {
                    if (Name == "IRDemoPass") {
                        outs() << "Registering IRDemoPass to run as described by pipeline\n";
                        FPM.addPass(IRDemoPlugin());
                        return true;
                    }
                    outs() << "Did not register IRDemoPass on function level even though plugin was loaded\n";
                    return false;
                });

            PB.registerPipelineParsingCallback(
                [](StringRef Name, ModulePassManager& MPM, ArrayRef<PassBuilder::PipelineElement>) {
                    if (Name == "require<MyAnalysis>") {
                        outs() << "Registering custom module level analysis as described by pipeline\n";
                        MPM.addPass(RequireAnalysisPass<MyAnalysis, Module>());
                        return true;
                    }
                    return false;
                });

            PB.registerAnalysisRegistrationCallback([](ModuleAnalysisManager& MAM) {
                MAM.registerPass([] {
                    return MyAnalysis();
                });
            });

            PB.registerAnalysisRegistrationCallback([](FunctionAnalysisManager& FAM) {
                FAM.registerPass([] {
                    return MyAnalysis();
                });
            });
        }
    };
}

extern "C" LLVM_ATTRIBUTE_WEAK PassPluginLibraryInfo llvmGetPassPluginInfo() { return getPluginInfo(); }