llvm/Analysis/InteractiveModelRunner.h

0001 //===- InteractiveModelRunner.h ---- "gym" ML model runner  -----*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008 //
0009
0010 #ifndef LLVM_ANALYSIS_INTERACTIVEMODELRUNNER_H
0011 #define LLVM_ANALYSIS_INTERACTIVEMODELRUNNER_H
0012
0013 #include "llvm/Analysis/MLModelRunner.h"
0014 #include "llvm/Analysis/TensorSpec.h"
0015 #include "llvm/Analysis/Utils/TrainingLogger.h"
0016 #include <system_error>
0017
0018 namespace llvm {
0019
0020 /// A MLModelRunner that asks for advice from an external agent, or host. It
0021 /// uses 2 files - ideally named pipes - one to send data to that agent, and
0022 /// one to receive advice.
0023 /// The data exchange uses the training logger (Utils/TrainingLogger.h) format.
0024 /// Specifically, the compiler will send the log header, set the context, and
0025 /// send observations; the host is expected to reply with a tensor value after
0026 /// each observation as a binary buffer that's conforming to the shape of the
0027 /// advice. Interleaved, the data closely resembles the training log for a
0028 /// log where we don't capture the reward signal.
0029 ///
0030 /// Note that the correctness of the received data is the responsibility of the
0031 /// host. In particular, if insufficient data were sent, the compiler will block
0032 /// when waiting for an advice.
0033 ///
0034 /// Note that the host can either open the pipes RW, or open first the pipe to
0035 /// the compiler - i.e. the "Inbound" - and then the "Outbound", to avoid
0036 /// deadlock. This is because the compiler first tries to open the inbound
0037 /// (which will hang until there's a writer on the other end).
0038 class InteractiveModelRunner : public MLModelRunner {
0039 public:
0040   InteractiveModelRunner(LLVMContext &Ctx,
0041                          const std::vector<TensorSpec> &Inputs,
0042                          const TensorSpec &Advice, StringRef OutboundName,
0043                          StringRef InboundName);
0044
0045   static bool classof(const MLModelRunner *R) {
0046     return R->getKind() == MLModelRunner::Kind::Interactive;
0047   }
0048   void switchContext(StringRef Name) override {
0049     Log->switchContext(Name);
0050     Log->flush();
0051   }
0052
0053   virtual ~InteractiveModelRunner();
0054
0055 private:
0056   void *evaluateUntyped() override;
0057   // This must be declared before InEC if we want to initialize it in the
0058   // ctor initializer list.
0059   int Inbound = -1;
0060   const std::vector<TensorSpec> InputSpecs;
0061   const TensorSpec OutputSpec;
0062   std::error_code OutEC;
0063   std::error_code InEC;
0064   std::vector<char> OutputBuffer;
0065   std::unique_ptr<Logger> Log;
0066 };
0067 } // namespace llvm
0068 #endif // LLVM_ANALYSIS_INTERACTIVEMODELRUNNER_H