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 MLIR lowering is enabled.
58/// 3. ReleaseModeModelRunner<CompiledModelType> otherwise.
59template <class CompiledModelType, bool HaveMLIRLowering, class CreateEmitCFunc>
60std::unique_ptr<MLModelRunner> createReleaseModeModelRunner(
61 LLVMContext &Ctx, const std::vector<TensorSpec> &InputFeatures,
62 StringRef DecisionName, const std::string &InteractiveChannelBaseName,
63 const TensorSpec &InteractiveDecisionSpec,
64 CreateEmitCFunc &&CreateEmitCModelRunner,
65 const EmbeddedModelRunnerOptions &Options = {}) {
66 if (!InteractiveChannelBaseName.empty()) {
67 return std::make_unique<InteractiveModelRunner>(
68 args&: Ctx, args: InputFeatures, args: InteractiveDecisionSpec,
69 args: InteractiveChannelBaseName + ".out",
70 args: InteractiveChannelBaseName + ".in");
71 }
72 if constexpr (HaveMLIRLowering) {
73 return CreateEmitCModelRunner(Ctx, InputFeatures);
74 } else {
75 return std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
76 Ctx, InputFeatures, DecisionName, Options);
77 }
78}
79
80} // namespace llvm
81
82#endif // LLVM_ANALYSIS_UTILS_MLGOUTILS_H
83