diff --git a/llvm/include/llvm/Analysis/InteractiveModelRunner.h b/llvm/include/llvm/Analysis/InteractiveModelRunner.h new file mode 100644 --- /dev/null +++ b/llvm/include/llvm/Analysis/InteractiveModelRunner.h @@ -0,0 +1,61 @@ +//===- InteractiveModelRunner.h ---- "gym" ML model runner -----*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// + +#ifndef LLVM_ANALYSIS_INTERACTIVEMODELRUNNER_H +#define LLVM_ANALYSIS_INTERACTIVEMODELRUNNER_H + +#include "llvm/Analysis/MLModelRunner.h" +#include "llvm/Analysis/TensorSpec.h" +#include "llvm/Analysis/Utils/TrainingLogger.h" +#include "llvm/Config/llvm-config.h" +#include "llvm/Support/raw_ostream.h" +#include + +namespace llvm { + +/// A MLModelRunner that asks for advice from an external agent, or host. It +/// uses 2 files - ideally named pipes - one to send data to that agent, and +/// one to receive advice. +/// The data exchange uses the training logger (Utils/TrainingLogger.h) format. +/// Specifically, the compiler will send the log header, set the context, and +/// send observations; the host is expected to reply with a tensor value after +/// each observation as a binary buffer that's conforming to the shape of the +/// advice. Interleaved, the data closely resembles the training log for a +/// log where we don't capture the reward signal. +/// +/// Note that the correctness of the received data is the responsibility of the +/// host. In particular, if insufficient data were sent, the compiler will block +/// when waiting for an advice. +class InteractiveModelRunner : public MLModelRunner { +public: + InteractiveModelRunner(LLVMContext &Ctx, + const std::vector &Inputs, + const TensorSpec &Advice, StringRef OutboundName, + StringRef InboundName); + + static bool classof(const MLModelRunner *R) { + return R->getKind() == MLModelRunner::Kind::Interactive; + } + void switchContext(StringRef Name) { + Log.switchContext(Name); + Log.flush(); + } + +private: + void *evaluateUntyped() override; + const std::vector InputSpecs; + const TensorSpec OutputSpec; + std::error_code OutEC; + std::error_code InEC; + raw_fd_stream Inbound; + std::vector OutputBuffer; + Logger Log; +}; +} // namespace llvm +#endif // LLVM_ANALYSIS_INTERACTIVEMODELRUNNER_H diff --git a/llvm/include/llvm/Analysis/MLModelRunner.h b/llvm/include/llvm/Analysis/MLModelRunner.h --- a/llvm/include/llvm/Analysis/MLModelRunner.h +++ b/llvm/include/llvm/Analysis/MLModelRunner.h @@ -47,7 +47,7 @@ return (const_cast(this))->getTensorUntyped(Index); } - enum class Kind : int { Unknown, Release, Development, NoOp }; + enum class Kind : int { Unknown, Release, Development, NoOp, Interactive }; Kind getKind() const { return Type; } protected: diff --git a/llvm/include/llvm/Analysis/TensorSpec.h b/llvm/include/llvm/Analysis/TensorSpec.h --- a/llvm/include/llvm/Analysis/TensorSpec.h +++ b/llvm/include/llvm/Analysis/TensorSpec.h @@ -103,6 +103,9 @@ size_t ElementSize = 0; }; +/// For debugging. +std::string tensorValueToString(const char *Buffer, const TensorSpec &Spec); + /// Construct a TensorSpec from a JSON dictionary of the form: /// { "name": , /// "port": , diff --git a/llvm/include/llvm/Analysis/Utils/TrainingLogger.h b/llvm/include/llvm/Analysis/Utils/TrainingLogger.h --- a/llvm/include/llvm/Analysis/Utils/TrainingLogger.h +++ b/llvm/include/llvm/Analysis/Utils/TrainingLogger.h @@ -116,6 +116,7 @@ void switchContext(StringRef Name); void startObservation(); void endObservation(); + void flush() { OS->flush(); } const std::string ¤tContext() const { return CurrentContext; } diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt --- a/llvm/lib/Analysis/CMakeLists.txt +++ b/llvm/lib/Analysis/CMakeLists.txt @@ -76,6 +76,7 @@ InstCount.cpp InstructionPrecedenceTracking.cpp InstructionSimplify.cpp + InteractiveModelRunner.cpp Interval.cpp IntervalPartition.cpp LazyBranchProbabilityInfo.cpp diff --git a/llvm/lib/Analysis/InteractiveModelRunner.cpp b/llvm/lib/Analysis/InteractiveModelRunner.cpp new file mode 100644 --- /dev/null +++ b/llvm/lib/Analysis/InteractiveModelRunner.cpp @@ -0,0 +1,77 @@ +//===- InteractiveModelRunner.cpp - noop ML model runner ----------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// A runner that communicates with an external agent via 2 file descriptors. +//===----------------------------------------------------------------------===// +#include "llvm/Analysis/InteractiveModelRunner.h" +#include "llvm/Analysis/MLModelRunner.h" +#include "llvm/Analysis/TensorSpec.h" +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/raw_ostream.h" + +using namespace llvm; + +#define _IMR_CL_VALS(T, N) clEnumValN(TensorType::N, #T, #T), + +static cl::opt DebugReply( + "interactive-model-runner-echo-type", cl::init(TensorType::Invalid), + cl::Hidden, + cl::desc("The InteractiveModelRunner will echo back to stderr " + "the data received " + "from the host as the specified type (for debugging purposes)."), + cl::values(SUPPORTED_TENSOR_TYPES(_IMR_CL_VALS) + clEnumValN(TensorType::Invalid, "disable", "Don't echo"))); + +#undef _IMR_CL_VALS + +InteractiveModelRunner::InteractiveModelRunner( + LLVMContext &Ctx, const std::vector &Inputs, + const TensorSpec &Advice, StringRef OutboundName, StringRef InboundName) + : MLModelRunner(Ctx, MLModelRunner::Kind::Interactive, Inputs.size()), + InputSpecs(Inputs), OutputSpec(Advice), Inbound(InboundName, InEC), + OutputBuffer(OutputSpec.getTotalTensorBufferSize()), + Log(std::make_unique(OutboundName, OutEC), InputSpecs, + Advice, /*IncludeReward=*/false) { + if (InEC) { + Ctx.emitError("Cannot open inbound file: " + InEC.message()); + return; + } + if (OutEC) { + Ctx.emitError("Cannot open outbound file: " + OutEC.message()); + return; + } + // Just like in the no inference case, this will allocate an appropriately + // sized buffer. + for (size_t I = 0; I < InputSpecs.size(); ++I) + setUpBufferForTensor(I, InputSpecs[I], nullptr); + Log.flush(); +} + +void *InteractiveModelRunner::evaluateUntyped() { + Log.startObservation(); + for (size_t I = 0; I < InputSpecs.size(); ++I) + Log.logTensorValue(I, reinterpret_cast(getTensorUntyped(I))); + Log.endObservation(); + Log.flush(); + + size_t InsPoint = 0; + char *Buff = OutputBuffer.data(); + const size_t Limit = OutputBuffer.size(); + while (InsPoint < Limit) { + auto Read = Inbound.read(Buff + InsPoint, OutputBuffer.size() - InsPoint); + if (Read < 0) { + Ctx.emitError("Failed reading from inbound file"); + break; + } + InsPoint += Read; + } + if (DebugReply != TensorType::Invalid) + dbgs() << tensorValueToString(OutputBuffer.data(), OutputSpec); + return OutputBuffer.data(); +} \ No newline at end of file diff --git a/llvm/lib/Analysis/TensorSpec.cpp b/llvm/lib/Analysis/TensorSpec.cpp --- a/llvm/lib/Analysis/TensorSpec.cpp +++ b/llvm/lib/Analysis/TensorSpec.cpp @@ -10,8 +10,10 @@ // utils. // //===----------------------------------------------------------------------===// +#include "llvm/ADT/STLExtras.h" #include "llvm/Config/config.h" +#include "llvm/ADT/StringExtras.h" #include "llvm/ADT/Twine.h" #include "llvm/Analysis/TensorSpec.h" #include "llvm/Support/CommandLine.h" @@ -102,4 +104,21 @@ return std::nullopt; } +std::string tensorValueToString(const char *Buffer, const TensorSpec &Spec) { + switch (Spec.type()) { +#define _IMR_DBG_PRINTER(T, N) \ + case TensorType::N: { \ + const T *TypedBuff = reinterpret_cast(Buffer); \ + auto R = llvm::make_range(TypedBuff, TypedBuff + Spec.getElementCount()); \ + return llvm::join( \ + llvm::map_range(R, [](T V) { return std::to_string(V); }), ","); \ + } + SUPPORTED_TENSOR_TYPES(_IMR_DBG_PRINTER) +#undef _IMR_DBG_PRINTER + case TensorType::Total: + case TensorType::Invalid: + llvm_unreachable("invalid tensor type"); + } +} + } // namespace llvm diff --git a/llvm/unittests/Analysis/MLModelRunnerTest.cpp b/llvm/unittests/Analysis/MLModelRunnerTest.cpp --- a/llvm/unittests/Analysis/MLModelRunnerTest.cpp +++ b/llvm/unittests/Analysis/MLModelRunnerTest.cpp @@ -7,10 +7,18 @@ //===----------------------------------------------------------------------===// #include "llvm/Analysis/MLModelRunner.h" +#include "llvm/Analysis/InteractiveModelRunner.h" #include "llvm/Analysis/NoInferenceModelRunner.h" #include "llvm/Analysis/ReleaseModeModelRunner.h" +#include "llvm/Support/BinaryByteStream.h" +#include "llvm/Support/FileUtilities.h" +#include "llvm/Support/JSON.h" +#include "llvm/Support/raw_ostream.h" #include "gtest/gtest.h" +#include +#include + using namespace llvm; namespace llvm { @@ -116,4 +124,135 @@ EXPECT_EQ(*Evaluator->getTensor(0), 1); EXPECT_EQ(*Evaluator->getTensor(1), 2); EXPECT_EQ(*Evaluator->getTensor(2), -3); +} + +TEST(InteractiveModelRunner, Evaluation) { + LLVMContext Ctx; + // Test the interaction with an external advisor by asking for advice twice. + // Use simple values, since we use the Logger underneath, that's tested more + // extensively elsewhere. + std::vector Inputs{ + TensorSpec::createSpec("a", {1}), + TensorSpec::createSpec("b", {1}), + TensorSpec::createSpec("c", {1}), + }; + TensorSpec AdviceSpec = TensorSpec::createSpec("advice", {1}); + + // Create the 2 files. Ideally we'd create them as named pipes, but that's not + // quite supported by the generic API. + std::error_code EC; + SmallString<64> FromCompilerName; + SmallString<64> ToCompilerName; + int FromCompilerFD = 0; + int ToCompilerFD = 0; + ASSERT_EQ(sys::fs::createTemporaryFile("InteractiveModelRunner_Evaluation", + "temp", FromCompilerFD, + FromCompilerName), + std::error_code()); + + ASSERT_EQ(sys::fs::createTemporaryFile("InteractiveModelRunner_Evaluation", + "temp", ToCompilerFD, ToCompilerName), + std::error_code()); + + raw_fd_stream FromCompiler(FromCompilerName, EC); + EXPECT_FALSE(EC); + raw_fd_ostream ToCompiler(ToCompilerName, EC); + EXPECT_FALSE(EC); + FileRemover Cleanup1(FromCompilerName); + FileRemover Cleanup2(ToCompilerName); + InteractiveModelRunner Evaluator(Ctx, Inputs, AdviceSpec, FromCompilerName, + ToCompilerName); + + Evaluator.switchContext("hi"); + + // Helper to read headers and other json lines. + SmallVector Buffer; + auto ReadLn = [&]() { + Buffer.clear(); + while (true) { + char Chr = 0; + auto Read = FromCompiler.read(&Chr, 1); + EXPECT_GE(Read, 0); + if (!Read) + continue; + if (Chr == '\n') + return StringRef(Buffer.data(), Buffer.size()); + Buffer.push_back(Chr); + } + }; + // See include/llvm/Analysis/Utils/TrainingLogger.h + // First comes the header + auto Header = json::parse(ReadLn()); + EXPECT_FALSE(Header.takeError()); + EXPECT_NE(Header->getAsObject()->getArray("features"), nullptr); + // Then comes the context + EXPECT_FALSE(json::parse(ReadLn()).takeError()); + + // Since the evaluator sends the features over and then blocks waiting for + // an answer, we must spawn a thread playing the role of the advisor / host: + std::atomic SeenObservations = 0; + std::thread Advisor([&]() { + EXPECT_EQ(SeenObservations, 0); + int64_t Features[3] = {0}; + auto FullyRead = [&]() { + size_t InsPt = 0; + const size_t ToRead = 3 * Inputs[0].getTotalTensorBufferSize(); + char *Buff = reinterpret_cast(Features); + while (InsPt < ToRead) { + auto Read = FromCompiler.read(Buff + InsPt, ToRead - InsPt); + EXPECT_GE(Read, 0); + InsPt += Read; + } + }; + // Observation + EXPECT_FALSE(json::parse(ReadLn()).takeError()); + // Tensor values + FullyRead(); + // a "\n" + char Chr = 0; + while (FromCompiler.read(&Chr, 1) == 0) { + } + EXPECT_EQ(Chr, '\n'); + EXPECT_EQ(Features[0], 42); + EXPECT_EQ(Features[1], 43); + EXPECT_EQ(Features[2], 100); + ++SeenObservations; + + // Send the advice + float Advice = 42.0012; + ToCompiler.write(reinterpret_cast(&Advice), + AdviceSpec.getTotalTensorBufferSize()); + ToCompiler.flush(); + + // Second observation, and same idea as above + EXPECT_FALSE(json::parse(ReadLn()).takeError()); + FullyRead(); + while (FromCompiler.read(&Chr, 1) == 0) { + } + EXPECT_EQ(Chr, '\n'); + EXPECT_EQ(Features[0], 10); + EXPECT_EQ(Features[1], -2); + EXPECT_EQ(Features[2], 1); + ++SeenObservations; + Advice = 50.30; + ToCompiler.write(reinterpret_cast(&Advice), + AdviceSpec.getTotalTensorBufferSize()); + ToCompiler.flush(); + }); + + EXPECT_EQ(SeenObservations, 0); + *Evaluator.getTensor(0) = 42; + *Evaluator.getTensor(1) = 43; + *Evaluator.getTensor(2) = 100; + float Ret = Evaluator.evaluate(); + EXPECT_EQ(SeenObservations, 1); + EXPECT_FLOAT_EQ(Ret, 42.0012); + + *Evaluator.getTensor(0) = 10; + *Evaluator.getTensor(1) = -2; + *Evaluator.getTensor(2) = 1; + Ret = Evaluator.evaluate(); + EXPECT_EQ(SeenObservations, 2); + EXPECT_FLOAT_EQ(Ret, 50.30); + Advisor.join(); } \ No newline at end of file diff --git a/llvm/unittests/Analysis/TensorSpecTest.cpp b/llvm/unittests/Analysis/TensorSpecTest.cpp --- a/llvm/unittests/Analysis/TensorSpecTest.cpp +++ b/llvm/unittests/Analysis/TensorSpecTest.cpp @@ -59,3 +59,10 @@ EXPECT_EQ(Spec3DLarge.getElementByteSize(), sizeof(float)); EXPECT_EQ(Spec1D.getElementByteSize(), sizeof(int16_t)); } + +TEST(TensorSpecTest, PrintValueForDebug) { + std::vector Values{1, 3}; + EXPECT_EQ(tensorValueToString(reinterpret_cast(Values.data()), + TensorSpec::createSpec("name", {2})), + "1,3"); +} \ No newline at end of file