path: root/contrib/libs/llvm12/include/llvm/Analysis/Utils/TFUtils.h

#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/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; 
}; 
 
/// 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: 
/// event 0: 
///   logTensorValue(0, ...) 
///   logTensorValue(1, ...) 
///   ... 
///   logReward(...) 
/// event 1: 
///   logTensorValue(0, ...) 
///   logTensorValue(1, ...) 
///   ... 
///   logReward(...) 
/// 
/// At the end, call print to generate the protobuf. 
class Logger final { 
public: 
  /// Construct a Logger. If IncludeReward is false, then logReward shouldn't 
  /// be called, and the reward feature won't be printed out. 
  Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs, 
         const TensorSpec &RewardSpec, bool IncludeReward) 
      : FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec), 
        RawLogData(FeatureSpecs.size() + IncludeReward), 
        IncludeReward(IncludeReward) {} 
 
  template <typename T> void logReward(T Value) { 
    assert(IncludeReward); 
    logTensorValue(RawLogData.size() - 1, &Value); 
  } 
 
  template <typename T> void logFinalReward(T Value) { 
    assert(RawLogData.back().empty()); 
    logReward(Value); 
  } 
 
  template <typename T> 
  void logTensorValue(size_t FeatureID, const T *Value, size_t Size = 1) { 
    const char *Start = reinterpret_cast<const char *>(Value); 
    const char *End = Start + sizeof(T) * Size; 
    RawLogData[FeatureID].insert(RawLogData[FeatureID].end(), Start, End); 
  } 
 
  void print(raw_ostream &OS); 
 
private: 
  std::vector<LoggedFeatureSpec> FeatureSpecs; 
  TensorSpec RewardSpec; 
  /// RawData has one entry per feature, plus one more for the reward. 
  /// Each feature's values are then stored in a vector, in succession. 
  /// This means the ith event is stored at [*][i] 
  std::vector<std::vector<char>> RawLogData; 
  const bool IncludeReward; 
}; 
 
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; }

private:
  void *getUntypedInput(size_t Index);
  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