aboutsummaryrefslogtreecommitdiffstats
path: root/contrib/libs/llvm14/include/llvm/Analysis/Utils/TFUtils.h
blob: de26f66b5750dfc8e98ec51a374042da7cdc7fcf (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
#pragma once

#ifdef __GNUC__
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-parameter"
#endif

//===- TFUtils.h - utilities for tensorflow C API ---------------*- 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_UTILS_TFUTILS_H
#define LLVM_ANALYSIS_UTILS_TFUTILS_H

#include "llvm/Config/llvm-config.h"

#ifdef LLVM_HAVE_TF_API
#include "llvm/ADT/StringMap.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/Support/JSON.h"

#include <memory>
#include <vector>

namespace llvm {

/// Load a SavedModel, find the given inputs and outputs, and setup storage
/// for input tensors. The user is responsible for correctly dimensioning the
/// input tensors and setting their values before calling evaluate().
/// To initialize:
/// - construct the object
/// - initialize the input tensors using initInput. Indices must correspond to
///   indices in the InputNames used at construction.
/// To use:
/// - set input values by using getInput to get each input tensor, and then
///   setting internal scalars, for all dimensions (tensors are row-major:
///   https://github.com/tensorflow/tensorflow/blob/r1.5/tensorflow/c/c_api.h#L205)
/// - call evaluate. The input tensors' values are not consumed after this, and
///   may still be read.
/// - use the outputs in the output vector
class TFModelEvaluatorImpl;
class EvaluationResultImpl;

/// TensorSpec encapsulates the specification of a tensor: its dimensions, or
/// "shape" (row-major), its type (see TensorSpec::getDataType specializations
/// for supported types), its name and port (see "TensorFlow: Large-Scale
/// Machine Learning on Heterogeneous Distributed Systems", section 4.2, para 2:
/// https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/45166.pdf)
///
/// TensorSpec is used to set up a TFModelEvaluator by describing the expected
/// inputs and outputs.
class TensorSpec final {
public:
  template <typename T>
  static TensorSpec createSpec(const std::string &Name,
                               const std::vector<int64_t> &Shape,
                               int Port = 0) {
    return TensorSpec(Name, Port, getDataType<T>(), Shape);
  }

  const std::string &name() const { return Name; }
  int port() const { return Port; }
  int typeIndex() const { return TypeIndex; }
  const std::vector<int64_t> &shape() const { return Shape; }

  bool operator==(const TensorSpec &Other) const {
    return Name == Other.Name && Port == Other.Port &&
           TypeIndex == Other.TypeIndex && Shape == Other.Shape;
  }

  bool operator!=(const TensorSpec &Other) const { return !(*this == Other); }

  /// Get the number of elements in a tensor with this shape.
  size_t getElementCount() const { return ElementCount; }
  /// Get the size, in bytes, of one element.
  size_t getElementByteSize() const;

  template <typename T> bool isElementType() const {
    return getDataType<T>() == TypeIndex;
  }

private:
  TensorSpec(const std::string &Name, int Port, int TypeIndex,
             const std::vector<int64_t> &Shape);

  template <typename T> static int getDataType() {
    llvm_unreachable("Undefined tensor type");
  }

  std::string Name;
  int Port = 0;
  int TypeIndex = 0;
  std::vector<int64_t> Shape;
  size_t ElementCount = 0;
};

/// Construct a TensorSpec from a JSON dictionary of the form:
/// { "name": <string>,
///   "port": <int>,
///   "type": <string. Use LLVM's types, e.g. float, double, int64_t>,
///   "shape": <array of ints> }
/// For the "type" field, see the C++ primitive types used in
/// TFUTILS_SUPPORTED_TYPES.
Optional<TensorSpec> getTensorSpecFromJSON(LLVMContext &Ctx,
                                           const json::Value &Value);

struct LoggedFeatureSpec {
  TensorSpec Spec;
  Optional<std::string> LoggingName;
  const std::string &getLoggingName() const {
    return LoggingName ? *LoggingName : Spec.name();
  }
};

/// Load the output specs. If SpecFileOverride is not empty, that path is used.
/// Otherwise, the file is assumed to be called 'output_spec.json' and be found
/// under ModelPath (the model directory).
/// The first output tensor name must match ExpectedDecisionName.
/// In case of error, the return is None and the error is logged.
Optional<std::vector<LoggedFeatureSpec>>
loadOutputSpecs(LLVMContext &Ctx, StringRef ExpectedDecisionName,
                StringRef ModelPath, StringRef SpecFileOverride = StringRef());

/// Logging utility - given an ordered specification of features, and assuming
/// a scalar reward, allow logging feature values and rewards, and then print
/// as tf.train.SequenceExample text protobuf.
/// The assumption is that, for an event to be logged (i.e. a set of feature
/// values and a reward), the user calls the log* API for each feature exactly
/// once, providing the index matching the position in the feature spec list
/// provided at construction. The example assumes the first feature's element
/// type is float, the second is int64, and the reward is float:
///
/// event 0:
///   logFloatValue(0, ...)
///   logInt64Value(1, ...)
///   ...
///   logFloatReward(...)
/// event 1:
///   logFloatValue(0, ...)
///   logInt64Value(1, ...)
///   ...
///   logFloatReward(...)
///
/// At the end, call print to generate the protobuf.
/// Alternatively, don't call logReward at the end of each event, just
/// log{Float|Int32|Int64}FinalReward at the end.
class LoggerDataImpl;
class Logger final {
public:
  /// Construct a Logger. If IncludeReward is false, then logReward or
  /// logFinalReward shouldn't be called, and the reward feature won't be
  /// printed out.
  /// NOTE: the FeatureSpecs are expected to be in the same order (i.e. have
  /// corresponding indices) with any MLModelRunner implementations
  /// corresponding to the model being trained/logged.
  Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs,
         const TensorSpec &RewardSpec, bool IncludeReward);

  ~Logger();

  void logFloatReward(float Value);
  void logInt32Reward(int32_t Value);
  void logInt64Reward(int64_t Value);

  void logFloatFinalReward(float Value);
  void logInt32FinalReward(int32_t Value);
  void logInt64FinalReward(int64_t Value);

  void logFloatValue(size_t FeatureID, const float *Value);
  void logInt32Value(size_t FeatureID, const int32_t *Value);
  void logInt64Value(size_t FeatureID, const int64_t *Value);

  void logSpecifiedTensorValue(size_t FeatureID, const char *RawData);

  // Warning! For int32_t, the return is set up for int64_t, so the caller needs
  // to piecemeal cast their int32_t values.
  // FIXME: let's drop int32_t support. While it's supported by evaluator, it's
  // not supported by the tensorflow::SequenceExample proto. For small values,
  // we can consider using bytes.
  char *addEntryAndGetFloatOrInt64Buffer(size_t FeatureID);

  // Flush the content of the log to the stream, clearing the stored data in the
  // process.
  void flush(std::string *Str);
  void flush(raw_ostream &OS);

  // Flush a set of logs that are produced from the same module, e.g.
  // per-function regalloc traces, as a google::protobuf::Struct message.
  static void flushLogs(raw_ostream &OS,
                        const StringMap<std::unique_ptr<Logger>> &Loggers);

private:
  std::vector<LoggedFeatureSpec> FeatureSpecs;
  TensorSpec RewardSpec;
  const bool IncludeReward;
  std::unique_ptr<LoggerDataImpl> LoggerData;
};

class TFModelEvaluator final {
public:
  /// The result of a model evaluation. Handles the lifetime of the output
  /// tensors, which means that their values need to be used before
  /// the EvaluationResult's dtor is called.
  class EvaluationResult {
  public:
    EvaluationResult(const EvaluationResult &) = delete;
    EvaluationResult &operator=(const EvaluationResult &Other) = delete;

    EvaluationResult(EvaluationResult &&Other);
    EvaluationResult &operator=(EvaluationResult &&Other);

    ~EvaluationResult();

    /// Get a (const) pointer to the first element of the tensor at Index.
    template <typename T> T *getTensorValue(size_t Index) {
      return static_cast<T *>(getUntypedTensorValue(Index));
    }

    template <typename T> const T *getTensorValue(size_t Index) const {
      return static_cast<T *>(getUntypedTensorValue(Index));
    }

    /// Get a (const) pointer to the untyped data of the tensor.
    void *getUntypedTensorValue(size_t Index);
    const void *getUntypedTensorValue(size_t Index) const;

  private:
    friend class TFModelEvaluator;
    EvaluationResult(std::unique_ptr<EvaluationResultImpl> Impl);
    std::unique_ptr<EvaluationResultImpl> Impl;
  };

  TFModelEvaluator(StringRef SavedModelPath,
                   const std::vector<TensorSpec> &InputSpecs,
                   const std::vector<TensorSpec> &OutputSpecs,
                   const char *Tags = "serve");
  TFModelEvaluator(StringRef SavedModelPath,
                   const std::vector<TensorSpec> &InputSpecs,
                   function_ref<TensorSpec(size_t)> GetOutputSpecs,
                   size_t OutputSpecsSize, const char *Tags = "serve");

  ~TFModelEvaluator();
  TFModelEvaluator(const TFModelEvaluator &) = delete;
  TFModelEvaluator(TFModelEvaluator &&) = delete;

  /// Evaluate the model, assuming it is valid. Returns None if the evaluation
  /// fails or the model is invalid, or an EvaluationResult otherwise. The
  /// inputs are assumed to have been already provided via getInput(). When
  /// returning None, it also invalidates this object.
  Optional<EvaluationResult> evaluate();

  /// Provides access to the input vector.
  template <typename T> T *getInput(size_t Index) {
    return static_cast<T *>(getUntypedInput(Index));
  }

  /// Returns true if the tensorflow model was loaded successfully, false
  /// otherwise.
  bool isValid() const { return !!Impl; }

  /// Untyped access to input.
  void *getUntypedInput(size_t Index);

private:
  std::unique_ptr<TFModelEvaluatorImpl> Impl;
};

/// List of supported types, as a pair:
/// - C++ type
/// - enum name (implementation-specific)
#define TFUTILS_SUPPORTED_TYPES(M)                                             \
  M(float, TF_FLOAT)                                                           \
  M(double, TF_DOUBLE)                                                         \
  M(int8_t, TF_INT8)                                                           \
  M(uint8_t, TF_UINT8)                                                         \
  M(int16_t, TF_INT16)                                                         \
  M(uint16_t, TF_UINT16)                                                       \
  M(int32_t, TF_INT32)                                                         \
  M(uint32_t, TF_UINT32)                                                       \
  M(int64_t, TF_INT64)                                                         \
  M(uint64_t, TF_UINT64)

#define TFUTILS_GETDATATYPE_DEF(T, E)                                          \
  template <> int TensorSpec::getDataType<T>();

TFUTILS_SUPPORTED_TYPES(TFUTILS_GETDATATYPE_DEF)

#undef TFUTILS_GETDATATYPE_DEF
} // namespace llvm

#endif // LLVM_HAVE_TF_API
#endif // LLVM_ANALYSIS_UTILS_TFUTILS_H

#ifdef __GNUC__
#pragma GCC diagnostic pop
#endif