1//===- MLGOUtils.h - Utilities for MLGO Release Mode ------------*- C++ -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8///
9/// \file
10/// This file provides helper functions for creating MLModelRunners and checking
11/// model validity in release mode.
12///
13//===----------------------------------------------------------------------===//
14
15#ifndef LLVM_ANALYSIS_UTILS_MLGOUTILS_H
16#define LLVM_ANALYSIS_UTILS_MLGOUTILS_H
17
18#include "llvm/ADT/StringRef.h"
19#include "llvm/Analysis/InteractiveModelRunner.h"
20#include "llvm/Analysis/MLModelRunner.h"
21#include "llvm/Analysis/ReleaseModeModelRunner.h"
22#include "llvm/Analysis/TensorSpec.h"
23#include "llvm/IR/LLVMContext.h"
24#include "llvm/Support/CommandLine.h"
25#include <memory>
26#include <string>
27#include <vector>
28
29namespace llvm {
30
31/// Helper to check if a release-mode ML advisor has a valid model to execute.
32/// Overload for cl::opt<EnumType>.
33template <class CompiledModelType, class EnumType, bool ExternalStorage,
34 class ParserClass>
35bool isReleaseModelValid(
36 StringRef InteractiveChannelBaseName,
37 const cl::opt<EnumType, ExternalStorage, ParserClass> &SelectedModel,
38 EnumType DefaultModelVal = EnumType::Default) {
39 return isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
40 !InteractiveChannelBaseName.empty() ||
41 SelectedModel != DefaultModelVal;
42}
43
44/// Helper to check if a release-mode ML advisor has a valid model to execute.
45/// Overload for plain EnumType.
46template <class CompiledModelType, class EnumType>
47bool isReleaseModelValid(StringRef InteractiveChannelBaseName,
48 EnumType SelectedModel,
49 EnumType DefaultModelVal = EnumType::Default) {
50 return isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
51 !InteractiveChannelBaseName.empty() ||
52 SelectedModel != DefaultModelVal;
53}
54
55/// Helper to construct the appropriate MLModelRunner in release mode:
56/// 1. InteractiveModelRunner if an interactive channel is specified.
57/// 2. EmitCModelRunner if an EmitC model is selected.
58/// 3. ReleaseModeModelRunner<CompiledModelType> if an embedded AOT model is
59/// available.
60template <class CompiledModelType, class CreateEmitCFunc>
61std::unique_ptr<MLModelRunner> createReleaseModeModelRunner(
62 LLVMContext &Ctx, const std::vector<TensorSpec> &InputFeatures,
63 StringRef DecisionName, const std::string &InteractiveChannelBaseName,
64 const TensorSpec &InteractiveDecisionSpec,
65 CreateEmitCFunc &&CreateEmitCModelRunner,
66 const EmbeddedModelRunnerOptions &Options = {}) {
67 if (!InteractiveChannelBaseName.empty()) {
68 return std::make_unique<InteractiveModelRunner>(
69 args&: Ctx, args: InputFeatures, args: InteractiveDecisionSpec,
70 args: InteractiveChannelBaseName + ".out",
71 args: InteractiveChannelBaseName + ".in");
72 }
73 if (auto EmitCRunner = CreateEmitCModelRunner(Ctx, InputFeatures))
74 return EmitCRunner;
75 if (isEmbeddedModelEvaluatorValid<CompiledModelType>()) {
76 auto Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
77 Ctx, InputFeatures, DecisionName, Options);
78 // Shapes the model was compiled for do not match the requested ones. The
79 // error is already reported, so let the caller fall back.
80 if (!Runner->isValid())
81 return nullptr;
82 return Runner;
83 }
84 return nullptr;
85}
86
87} // namespace llvm
88
89#endif // LLVM_ANALYSIS_UTILS_MLGOUTILS_H
90