[llvm] 62fc44c - [MLInliner] In development mode, obtain the output specs from a file
Mircea Trofin via llvm-commits
llvm-commits at lists.llvm.org
Mon Aug 17 16:56:57 PDT 2020
Author: Mircea Trofin
Date: 2020-08-17T16:56:47-07:00
New Revision: 62fc44ca3cf66442b30e22b1be34afc492a2a388
URL: https://github.com/llvm/llvm-project/commit/62fc44ca3cf66442b30e22b1be34afc492a2a388
DIFF: https://github.com/llvm/llvm-project/commit/62fc44ca3cf66442b30e22b1be34afc492a2a388.diff
LOG: [MLInliner] In development mode, obtain the output specs from a file
Different training algorithms may produce models that, besides the main
policy output (i.e. inline/don't inline), produce additional outputs
that are necessary for the next training stage. To facilitate this, in
development mode, we require the training policy infrastructure produce
a description of the outputs that are interesting to it, in the form of
a JSON file. We special-case the first entry in the JSON file as the
inlining decision - we care about its value, so we can guide inlining
during training - but treat the rest as opaque data that we just copy
over to the training log.
Differential Revision: https://reviews.llvm.org/D85674
Added:
llvm/lib/Analysis/models/inliner/output_spec.json
llvm/test/Transforms/Inline/ML/Inputs/test_output_spec.json
Modified:
llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp
llvm/test/Transforms/Inline/ML/development-training-log.ll
Removed:
################################################################################
diff --git a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp
index b20f9427525a..19ec8cbd4042 100644
--- a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp
@@ -21,6 +21,7 @@
#include "llvm/IR/LLVMContext.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/ManagedStatic.h"
+#include "llvm/Support/Path.h"
#include <vector>
@@ -32,17 +33,43 @@ static cl::opt<std::string> TrainingLog(
static cl::opt<std::string> TFModelUnderTrainingPath(
"ml-inliner-model-under-training", cl::Hidden,
- cl::desc("Path to SavedModel from the previous training iteration."));
+ cl::desc(R"(Path to SavedModel from the previous training iteration.
+The directory is also expected to contain a JSON specification of the
+outputs expected to be logged, where the first entry must be the
+inlining decision. The file containing the specification should be
+called output_spec.json. The expected JSON value is an array of
+dictionaries. Each dictionary should have 2 keys:
+
+- "tensor_spec, followed by the TensorSpec description of the
+output; and
+- "logging_name", a string indicating the name to use when
+logging the output values.
+
+Example:
+[
+ {
+ "logging_name" : "some_name",
+ "tensor_spec" : {
+ "name" : "model_name",
+ "port" : 0,
+ "shape" : [2, 3],
+ "type" : "float"
+ }
+ }
+]
+
+The first value must always correspond to the decision.)"));
+
+static cl::opt<std::string> TFOutputSpecOverride(
+ "ml-inliner-output-spec-override", cl::Hidden,
+ cl::desc("Override the path to the output spec json file. See "
+ "-ml-inliner-model-under-training documentation for the "
+ "specification of that file."));
static cl::opt<std::string> TFFeedPrefix("ml-inliner-trained-model-feed-prefix",
cl::Hidden, cl::init("action_"),
cl::desc("Prefix for feature names."));
-static cl::opt<std::string> TFDecisionName(
- "ml-inliner-trained-model-decision-name", cl::Hidden,
- cl::init("StatefulPartitionedCall"),
- cl::desc("Name of the graph operation representing the decision."));
-
namespace {
/// An InlineEvent, used by TrainingLogger.
struct InlineEvent {
@@ -69,9 +96,10 @@ struct InlineEvent {
/// Because this is a protobuf, we cannot just stream the events as they come.
/// Internally, TrainingLogger stores data in column-major format, because that
/// lines up with how TF SequenceExample represents it.
+class ModelUnderTrainingRunner;
class TrainingLogger final {
public:
- TrainingLogger(StringRef LogFileName);
+ TrainingLogger(StringRef LogFileName, const ModelUnderTrainingRunner *MUTR);
/// Log one inlining event.
void logInlineEvent(const InlineEvent &Event,
@@ -157,9 +185,13 @@ class TrainingLogger final {
}
StringRef LogFileName;
+ const ModelUnderTrainingRunner *const MUTR;
std::vector<InlineFeatures> Features;
std::vector<int64_t> DefaultDecisions;
- std::vector<int64_t> Decisions;
+ // We store all outputs as data blobs, but we always expect to have one, the
+ // first one, representing the decision. While we could track that separately,
+ // for uniformity, we store it, generically, here.
+ std::vector<std::vector<char>> Outputs;
std::vector<bool> Effects;
std::vector<int64_t> Rewards;
};
@@ -336,8 +368,22 @@ class ModelUnderTrainingRunner final : public MLModelRunner {
int64_t getFeature(int Index) const override;
bool isValid() const { return !!Evaluator; }
+ const std::vector<std::string> outputNames() const { return OutputNames; }
+
+ const std::vector<TensorSpec> outputSpecs() const { return OutputSpecs; }
+
+ const Optional<TFModelEvaluator::EvaluationResult> &
+ lastEvaluationResult() const {
+ return LastEvaluationResult;
+ }
+
private:
std::unique_ptr<TFModelEvaluator> Evaluator;
+ std::vector<std::string> OutputNames;
+ std::vector<TensorSpec> OutputSpecs;
+ Optional<TFModelEvaluator::EvaluationResult> LastEvaluationResult;
+
+ bool loadOutputSpecs(LLVMContext &Ctx, StringRef FileName);
// The training framework needs some additional features.
const std::vector<TensorSpec> TrainingOnlyFeatures{
@@ -348,10 +394,15 @@ class ModelUnderTrainingRunner final : public MLModelRunner {
};
} // namespace
-TrainingLogger::TrainingLogger(StringRef LogFileName)
- : LogFileName(LogFileName) {
+TrainingLogger::TrainingLogger(StringRef LogFileName,
+ const ModelUnderTrainingRunner *MUTR)
+ : LogFileName(LogFileName), MUTR(MUTR) {
for (size_t I = 0; I < NumberOfFeatures; ++I)
Features.push_back(InlineFeatures());
+
+ // The first output is the inlining decision.
+ auto OutputCount = MUTR ? MUTR->outputSpecs().size() : 1;
+ Outputs.assign(OutputCount, std::vector<char>());
}
/// Log one inlining event.
@@ -360,16 +411,27 @@ void TrainingLogger::logInlineEvent(const InlineEvent &Event,
for (size_t I = 0; I < NumberOfFeatures; ++I)
Features[I].push_back(ModelRunner.getFeature(I));
- Decisions.push_back(Event.AdvisedDecision);
Effects.push_back(Event.Effect);
Rewards.push_back(Event.Reward);
DefaultDecisions.push_back(Event.DefaultDecision);
+ int64_t Advice = static_cast<int64_t>(Event.AdvisedDecision);
+ const char *AdviceData = reinterpret_cast<const char *>(&Advice);
+ Outputs[0].insert(Outputs[0].end(), AdviceData, AdviceData + sizeof(int64_t));
+ for (size_t I = 1; I < Outputs.size(); ++I) {
+ const auto &Result = *MUTR->lastEvaluationResult();
+ auto &Spec = MUTR->outputSpecs()[I];
+ const char *RawData =
+ reinterpret_cast<const char *>(Result.getUntypedTensorValue(I));
+ Outputs[I].insert(Outputs[I].end(), RawData,
+ RawData +
+ Spec.getElementCount() * Spec.getElementByteSize());
+ }
}
void TrainingLogger::print() {
std::error_code EC;
raw_fd_ostream OutFile(LogFileName, EC);
- size_t NumberOfRecords = Decisions.size();
+ size_t NumberOfRecords = Rewards.size();
if (NumberOfRecords == 0)
return;
@@ -383,13 +445,18 @@ void TrainingLogger::print() {
OutFile, TensorSpec::createSpec<int64_t>(DefaultDecisionName, {1}),
DefaultDecisions.data(), NumberOfRecords);
- writeTensorsAsFeatureLists(OutFile,
- TensorSpec::createSpec<int64_t>(DecisionName, {1}),
- Decisions.data(), NumberOfRecords);
+ writeRawTensorsAsFeatureLists(
+ OutFile, TensorSpec::createSpec<int64_t>(DecisionName, {1}),
+ Outputs[0].data(), NumberOfRecords);
writeTensorsAsFeatureLists(OutFile,
TensorSpec::createSpec<int64_t>(RewardName, {1}),
Rewards.data(), NumberOfRecords);
+ for (size_t I = 1; I < Outputs.size(); ++I)
+ writeRawTensorsAsFeatureLists(OutFile, MUTR->outputSpecs()[I],
+ Outputs[I].data(), NumberOfRecords,
+ StringRef(MUTR->outputNames()[I]));
+
OutFile << "}\n";
}
@@ -472,13 +539,19 @@ ModelUnderTrainingRunner::ModelUnderTrainingRunner(LLVMContext &Ctx,
const std::string &ModelPath)
: MLModelRunner(Ctx) {
std::vector<TensorSpec> InputSpecs;
- std::vector<TensorSpec> OutputSpecs;
for (size_t I = 0; I < NumberOfFeatures; ++I)
InputSpecs.push_back(
TensorSpec::createSpec<int64_t>(TFFeedPrefix + FeatureNameMap[I], {1}));
InputSpecs.insert(InputSpecs.end(), TrainingOnlyFeatures.begin(),
TrainingOnlyFeatures.end());
- OutputSpecs.push_back(TensorSpec::createSpec<int64_t>(TFDecisionName, {1}));
+ SmallVector<char, 128> OutputSpecsPath;
+ StringRef OutputSpecPath = TFOutputSpecOverride;
+ if (OutputSpecPath.empty()) {
+ llvm::sys::path::append(OutputSpecsPath, ModelPath, "output_spec.json");
+ OutputSpecPath = {OutputSpecsPath.data(), OutputSpecsPath.size()};
+ }
+ if (!loadOutputSpecs(Ctx, OutputSpecPath))
+ return;
Evaluator =
std::make_unique<TFModelEvaluator>(ModelPath, InputSpecs, OutputSpecs);
@@ -489,13 +562,70 @@ ModelUnderTrainingRunner::ModelUnderTrainingRunner(LLVMContext &Ctx,
}
}
+bool ModelUnderTrainingRunner::loadOutputSpecs(LLVMContext &Ctx,
+ StringRef FileName) {
+ auto BufferOrError = MemoryBuffer::getFileOrSTDIN(FileName);
+ if (!BufferOrError) {
+ Ctx.emitError("Error opening output specs file: " + FileName + " : " +
+ BufferOrError.getError().message());
+ return false;
+ }
+ auto ParsedJSONValues = json::parse(BufferOrError.get()->getBuffer());
+ if (!ParsedJSONValues) {
+ Ctx.emitError("Could not parse specs file: " + FileName);
+ return false;
+ }
+ auto ValuesArray = ParsedJSONValues->getAsArray();
+ if (!ValuesArray) {
+ Ctx.emitError("Expected an array of {tensor_spec:<TensorSpec>, "
+ "logging_name:<name>} dictionaries");
+ return false;
+ }
+
+ for (const auto &Value : *ValuesArray)
+ if (const auto *Obj = Value.getAsObject())
+ if (const auto *SpecPart = Obj->get("tensor_spec"))
+ if (auto TensorSpec = getTensorSpecFromJSON(Ctx, *SpecPart))
+ if (auto LoggingName = Obj->getString("logging_name")) {
+ if (!TensorSpec->isElementType<int64_t>() &&
+ !TensorSpec->isElementType<int32_t>() &&
+ !TensorSpec->isElementType<float>()) {
+ Ctx.emitError(
+ "Only int64, int32, and float tensors are supported. "
+ "Found unsupported type for tensor named " +
+ TensorSpec->name());
+ return false;
+ }
+ OutputNames.push_back(LoggingName->str());
+ OutputSpecs.push_back(*TensorSpec);
+ }
+
+ if (ValuesArray->size() != OutputNames.size()) {
+ Ctx.emitError(
+ "Unable to parse output spec. It should be a json file containing an "
+ "array of dictionaries. Each dictionary must have a 'tensor_spec' key, "
+ "with a json object describing a TensorSpec; and a 'logging_name' key, "
+ "which is a string to use as name when logging this tensor in the "
+ "training log.");
+ return false;
+ }
+ assert(OutputNames.size() == OutputSpecs.size());
+ if (OutputNames.empty() || OutputNames[0] != DecisionName) {
+ Ctx.emitError("The first output spec must describe the decision tensor, "
+ "and must have the logging_name " +
+ StringRef(DecisionName));
+ return false;
+ }
+ return true;
+}
+
bool ModelUnderTrainingRunner::run() {
- auto ER = Evaluator->evaluate();
- if (!ER.hasValue()) {
+ LastEvaluationResult = Evaluator->evaluate();
+ if (!LastEvaluationResult.hasValue()) {
Ctx.emitError("Error evaluating model.");
return false;
}
- int64_t Decision = *ER->getTensorValue<int64_t>(0);
+ int64_t Decision = *LastEvaluationResult->getTensorValue<int64_t>(0);
return static_cast<bool>(Decision);
}
@@ -521,22 +651,24 @@ std::unique_ptr<InlineAdvisor> llvm::getDevelopmentModeAdvisor(
}
std::unique_ptr<MLModelRunner> Runner;
-
+ ModelUnderTrainingRunner *MUTRPtr = nullptr;
bool IsDoingInference = false;
if (TFModelUnderTrainingPath.empty())
Runner.reset(new NoInferenceModelRunner(Ctx));
else {
- Runner = std::make_unique<ModelUnderTrainingRunner>(
+ auto MUTR = std::make_unique<ModelUnderTrainingRunner>(
Ctx, TFModelUnderTrainingPath);
- if (!Runner) {
+ if (!MUTR || !MUTR->isValid()) {
Ctx.emitError("Could not load the policy model from the provided path");
return nullptr;
}
IsDoingInference = true;
+ MUTRPtr = MUTR.get();
+ Runner = std::move(MUTR);
}
std::unique_ptr<TrainingLogger> Logger;
if (!TrainingLog.empty())
- Logger = std::make_unique<TrainingLogger>(TrainingLog);
+ Logger = std::make_unique<TrainingLogger>(TrainingLog, MUTRPtr);
return std::make_unique<DevelopmentModeMLInlineAdvisor>(
M, MAM, std::move(Runner), GetDefaultAdvice, IsDoingInference,
diff --git a/llvm/lib/Analysis/models/inliner/output_spec.json b/llvm/lib/Analysis/models/inliner/output_spec.json
new file mode 100644
index 000000000000..d9e2060cf176
--- /dev/null
+++ b/llvm/lib/Analysis/models/inliner/output_spec.json
@@ -0,0 +1,14 @@
+[
+ {
+ "logging_name": "inlining_decision",
+ "tensor_spec": {
+ "name": "StatefulPartitionedCall",
+ "port": 0,
+ "type": "int64",
+ "shape": [
+ 1
+ ]
+ }
+ }
+]
+
\ No newline at end of file
diff --git a/llvm/test/Transforms/Inline/ML/Inputs/test_output_spec.json b/llvm/test/Transforms/Inline/ML/Inputs/test_output_spec.json
new file mode 100644
index 000000000000..bd6a19c9572b
--- /dev/null
+++ b/llvm/test/Transforms/Inline/ML/Inputs/test_output_spec.json
@@ -0,0 +1,25 @@
+[
+ {
+ "logging_name": "inlining_decision",
+ "tensor_spec": {
+ "name": "StatefulPartitionedCall",
+ "port": 0,
+ "type": "int64",
+ "shape": [
+ 1
+ ]
+ }
+ },
+ {
+ "logging_name": "fake_extra_output",
+ "tensor_spec": {
+ "name": "StatefulPartitionedCall",
+ "port": 0,
+ "type": "int64",
+ "shape": [
+ 1
+ ]
+ }
+ }
+]
+
\ No newline at end of file
diff --git a/llvm/test/Transforms/Inline/ML/development-training-log.ll b/llvm/test/Transforms/Inline/ML/development-training-log.ll
index 4bf6259f7e48..77347455958b 100644
--- a/llvm/test/Transforms/Inline/ML/development-training-log.ll
+++ b/llvm/test/Transforms/Inline/ML/development-training-log.ll
@@ -1,6 +1,7 @@
; Test that we can produce a log if we have or do not have a model, in development mode.
; REQUIRES: have_tf_api
; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -ml-inliner-model-under-training=%S/../../../../lib/Analysis/models/inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -S < %s | FileCheck %s
+; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -ml-inliner-model-under-training=%S/../../../../lib/Analysis/models/inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-output-spec-override=%S/Inputs/test_output_spec.json -S < %s | FileCheck %s --check-prefix=EXTRA-OUTPUTS
; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -S < %s | FileCheck %s
target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
@@ -48,4 +49,9 @@ define dso_local i32 @top() {
; CHECK-NEXT: key: "delta_size" value: {
; CHECK-NEXT: feature: { int64_list: { value: [0] } }
; CHECK-NEXT: }
-; CHECK-NEXT: }
\ No newline at end of file
+; CHECK-NEXT: }
+; CHECK-NOT: fake_extra_output
+; EXTRA-OUTPUTS: key: "fake_extra_output" value: {
+; EXTRA-OUTPUTS-NEXT: feature: { int64_list: { value: [1] } }
+; EXTRA-OUTPUTS-NEXT: }
+; EXTRA-OUTPUTS-NEXT: }
\ No newline at end of file
More information about the llvm-commits
mailing list