xref: /llvm-project/llvm/include/llvm/Analysis/InteractiveModelRunner.h (revision 0da2ba811ac8a01509bc533428941fb9519c0715)
1 //===- InteractiveModelRunner.h ---- "gym" ML model runner  -----*- 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 
10 #ifndef LLVM_ANALYSIS_INTERACTIVEMODELRUNNER_H
11 #define LLVM_ANALYSIS_INTERACTIVEMODELRUNNER_H
12 
13 #include "llvm/Analysis/MLModelRunner.h"
14 #include "llvm/Analysis/TensorSpec.h"
15 #include "llvm/Analysis/Utils/TrainingLogger.h"
16 #include <system_error>
17 
18 namespace llvm {
19 
20 /// A MLModelRunner that asks for advice from an external agent, or host. It
21 /// uses 2 files - ideally named pipes - one to send data to that agent, and
22 /// one to receive advice.
23 /// The data exchange uses the training logger (Utils/TrainingLogger.h) format.
24 /// Specifically, the compiler will send the log header, set the context, and
25 /// send observations; the host is expected to reply with a tensor value after
26 /// each observation as a binary buffer that's conforming to the shape of the
27 /// advice. Interleaved, the data closely resembles the training log for a
28 /// log where we don't capture the reward signal.
29 ///
30 /// Note that the correctness of the received data is the responsibility of the
31 /// host. In particular, if insufficient data were sent, the compiler will block
32 /// when waiting for an advice.
33 ///
34 /// Note that the host can either open the pipes RW, or open first the pipe to
35 /// the compiler - i.e. the "Inbound" - and then the "Outbound", to avoid
36 /// deadlock. This is because the compiler first tries to open the inbound
37 /// (which will hang until there's a writer on the other end).
38 class InteractiveModelRunner : public MLModelRunner {
39 public:
40   InteractiveModelRunner(LLVMContext &Ctx,
41                          const std::vector<TensorSpec> &Inputs,
42                          const TensorSpec &Advice, StringRef OutboundName,
43                          StringRef InboundName);
44 
45   static bool classof(const MLModelRunner *R) {
46     return R->getKind() == MLModelRunner::Kind::Interactive;
47   }
48   void switchContext(StringRef Name) override {
49     Log->switchContext(Name);
50     Log->flush();
51   }
52 
53   virtual ~InteractiveModelRunner();
54 
55 private:
56   void *evaluateUntyped() override;
57   // This must be declared before InEC if we want to initialize it in the
58   // ctor initializer list.
59   int Inbound = -1;
60   const std::vector<TensorSpec> InputSpecs;
61   const TensorSpec OutputSpec;
62   std::error_code OutEC;
63   std::error_code InEC;
64   std::vector<char> OutputBuffer;
65   std::unique_ptr<Logger> Log;
66 };
67 } // namespace llvm
68 #endif // LLVM_ANALYSIS_INTERACTIVEMODELRUNNER_H
69