//===- 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/config.h" #ifdef LLVM_HAVE_TF_API #include "llvm/IR/LLVMContext.h" #include #include 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; 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(EvaluationResult &&Other); ~EvaluationResult(); /// Get a pointer to the first element of the tensor at Index. template T *getTensorValue(size_t Index) { return static_cast(getUntypedTensorValue(Index)); } private: friend class TFModelEvaluator; EvaluationResult(std::unique_ptr Impl); void *getUntypedTensorValue(size_t Index); std::unique_ptr Impl; }; TFModelEvaluator(StringRef SavedModelPath, const std::vector &InputNames, const std::vector &OutputNames, 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 evaluate(); /// Provides access to the input vector. template T *getInput(size_t Index) { return static_cast(getUntypedInput(Index)); } /// Returns true if the tensorflow model was loaded successfully, false /// otherwise. bool isValid() const { return !!Impl; } /// Initialize the input at Index as a tensor of the given type and /// dimensions. template void initInput(size_t Index, const std::vector &Dimensions) { return initInput(Index, getModelTypeIndex(), Dimensions); } private: void *getUntypedInput(size_t Index); template int getModelTypeIndex(); void initInput(size_t Index, int TypeIndex, const std::vector &Dimensions); std::unique_ptr Impl; }; template <> int TFModelEvaluator::getModelTypeIndex(); template <> int TFModelEvaluator::getModelTypeIndex(); template <> int TFModelEvaluator::getModelTypeIndex(); template <> int TFModelEvaluator::getModelTypeIndex(); template <> int TFModelEvaluator::getModelTypeIndex(); template <> int TFModelEvaluator::getModelTypeIndex(); template <> int TFModelEvaluator::getModelTypeIndex(); template <> int TFModelEvaluator::getModelTypeIndex(); template <> int TFModelEvaluator::getModelTypeIndex(); template <> int TFModelEvaluator::getModelTypeIndex(); } // namespace llvm #endif // LLVM_HAVE_TF_API #endif // LLVM_ANALYSIS_UTILS_TFUTILS_H