1 //===- TensorSpec.cpp - tensor type abstraction ---------------------------===//
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 // Implementation file for the abstraction of a tensor type, and JSON loading
10 // utils.
11 //
12 //===----------------------------------------------------------------------===//
13 #include "llvm/Config/config.h"
14
15 #include "llvm/ADT/Twine.h"
16 #include "llvm/Analysis/TensorSpec.h"
17 #include "llvm/Support/CommandLine.h"
18 #include "llvm/Support/Debug.h"
19 #include "llvm/Support/JSON.h"
20 #include "llvm/Support/ManagedStatic.h"
21 #include "llvm/Support/MemoryBuffer.h"
22 #include "llvm/Support/Path.h"
23 #include "llvm/Support/raw_ostream.h"
24 #include <cassert>
25 #include <numeric>
26
27 using namespace llvm;
28
29 namespace llvm {
30
31 #define TFUTILS_GETDATATYPE_IMPL(T, E) \
32 template <> TensorType TensorSpec::getDataType<T>() { return TensorType::E; }
33
SUPPORTED_TENSOR_TYPES(TFUTILS_GETDATATYPE_IMPL)34 SUPPORTED_TENSOR_TYPES(TFUTILS_GETDATATYPE_IMPL)
35
36 #undef TFUTILS_GETDATATYPE_IMPL
37
38 TensorSpec::TensorSpec(const std::string &Name, int Port, TensorType Type,
39 size_t ElementSize, const std::vector<int64_t> &Shape)
40 : Name(Name), Port(Port), Type(Type), Shape(Shape),
41 ElementCount(std::accumulate(Shape.begin(), Shape.end(), 1,
42 std::multiplies<int64_t>())),
43 ElementSize(ElementSize) {}
44
getTensorSpecFromJSON(LLVMContext & Ctx,const json::Value & Value)45 Optional<TensorSpec> getTensorSpecFromJSON(LLVMContext &Ctx,
46 const json::Value &Value) {
47 auto EmitError = [&](const llvm::Twine &Message) -> Optional<TensorSpec> {
48 std::string S;
49 llvm::raw_string_ostream OS(S);
50 OS << Value;
51 Ctx.emitError("Unable to parse JSON Value as spec (" + Message + "): " + S);
52 return None;
53 };
54 // FIXME: accept a Path as a parameter, and use it for error reporting.
55 json::Path::Root Root("tensor_spec");
56 json::ObjectMapper Mapper(Value, Root);
57 if (!Mapper)
58 return EmitError("Value is not a dict");
59
60 std::string TensorName;
61 int TensorPort = -1;
62 std::string TensorType;
63 std::vector<int64_t> TensorShape;
64
65 if (!Mapper.map<std::string>("name", TensorName))
66 return EmitError("'name' property not present or not a string");
67 if (!Mapper.map<std::string>("type", TensorType))
68 return EmitError("'type' property not present or not a string");
69 if (!Mapper.map<int>("port", TensorPort))
70 return EmitError("'port' property not present or not an int");
71 if (!Mapper.map<std::vector<int64_t>>("shape", TensorShape))
72 return EmitError("'shape' property not present or not an int array");
73
74 #define PARSE_TYPE(T, E) \
75 if (TensorType == #T) \
76 return TensorSpec::createSpec<T>(TensorName, TensorShape, TensorPort);
77 SUPPORTED_TENSOR_TYPES(PARSE_TYPE)
78 #undef PARSE_TYPE
79 return None;
80 }
81
82 Optional<std::vector<LoggedFeatureSpec>>
loadOutputSpecs(LLVMContext & Ctx,StringRef ExpectedDecisionName,StringRef ModelPath,StringRef SpecFileOverride)83 loadOutputSpecs(LLVMContext &Ctx, StringRef ExpectedDecisionName,
84 StringRef ModelPath, StringRef SpecFileOverride) {
85 SmallVector<char, 128> OutputSpecsPath;
86 StringRef FileName = SpecFileOverride;
87 if (FileName.empty()) {
88 llvm::sys::path::append(OutputSpecsPath, ModelPath, "output_spec.json");
89 FileName = {OutputSpecsPath.data(), OutputSpecsPath.size()};
90 }
91
92 auto BufferOrError = MemoryBuffer::getFileOrSTDIN(FileName);
93 if (!BufferOrError) {
94 Ctx.emitError("Error opening output specs file: " + FileName + " : " +
95 BufferOrError.getError().message());
96 return None;
97 }
98 auto ParsedJSONValues = json::parse(BufferOrError.get()->getBuffer());
99 if (!ParsedJSONValues) {
100 Ctx.emitError("Could not parse specs file: " + FileName);
101 return None;
102 }
103 auto ValuesArray = ParsedJSONValues->getAsArray();
104 if (!ValuesArray) {
105 Ctx.emitError("Expected an array of {tensor_spec:<TensorSpec>, "
106 "logging_name:<name>} dictionaries");
107 return None;
108 }
109 std::vector<LoggedFeatureSpec> Ret;
110 for (const auto &Value : *ValuesArray)
111 if (const auto *Obj = Value.getAsObject())
112 if (const auto *SpecPart = Obj->get("tensor_spec"))
113 if (auto TensorSpec = getTensorSpecFromJSON(Ctx, *SpecPart))
114 if (auto LoggingName = Obj->getString("logging_name")) {
115 if (!TensorSpec->isElementType<int64_t>() &&
116 !TensorSpec->isElementType<int32_t>() &&
117 !TensorSpec->isElementType<float>()) {
118 Ctx.emitError(
119 "Only int64, int32, and float tensors are supported. "
120 "Found unsupported type for tensor named " +
121 TensorSpec->name());
122 return None;
123 }
124 Ret.push_back({*TensorSpec, LoggingName->str()});
125 }
126
127 if (ValuesArray->size() != Ret.size()) {
128 Ctx.emitError(
129 "Unable to parse output spec. It should be a json file containing an "
130 "array of dictionaries. Each dictionary must have a 'tensor_spec' key, "
131 "with a json object describing a TensorSpec; and a 'logging_name' key, "
132 "which is a string to use as name when logging this tensor in the "
133 "training log.");
134 return None;
135 }
136 if (Ret.empty() || *Ret[0].LoggingName != ExpectedDecisionName) {
137 Ctx.emitError("The first output spec must describe the decision tensor, "
138 "and must have the logging_name " +
139 StringRef(ExpectedDecisionName));
140 return None;
141 }
142 return Ret;
143 }
144 } // namespace llvm
145