TensorSpec.h

Go to the documentation of this file.

//===- TensorSpec.h - type descriptor for a tensor --------------*- C++ -*-===//
//
// Part of the MLCompilerBridge Project, under the Apache License v2.0 with LLVM
// Exceptions. See the LICENSE file for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
// (Preliminary version adopted from TensorSpec.h of LLVM 17.X)
//
//===----------------------------------------------------------------------===//
 
#ifndef MLBRIDGE_TENSORSPEC_H
#define MLBRIDGE_TENSORSPEC_H
 
#include "llvm/ADT/Optional.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/Support/JSON.h"
#include <memory>
#include <numeric>
#include <vector>
 
namespace MLBridge {
#define SUPPORTED_TENSOR_TYPES(M)                                              \
  M(float, Float)                                                              \
  M(double, Double)                                                            \
  M(int8_t, Int8)                                                              \
  M(uint8_t, UInt8)                                                            \
  M(int16_t, Int16)                                                            \
  M(uint16_t, UInt16)                                                          \
  M(int32_t, Int32)                                                            \
  M(uint32_t, UInt32)                                                          \
  M(int64_t, Int64)                                                            \
  M(uint64_t, UInt64)
 
enum class TensorType {
  Invalid,
#define _TENSOR_TYPE_ENUM_MEMBERS(_, Name) Name,
  SUPPORTED_TENSOR_TYPES(_TENSOR_TYPE_ENUM_MEMBERS)
#undef _TENSOR_TYPE_ENUM_MEMBERS
      Total
};
 
class TensorSpec final {
public:
  template <typename T>
  static TensorSpec createSpec(const std::string &Name,
                               const std::vector<int64_t> &Shape,
                               int Port = 6020) {
    return TensorSpec(Name, Port, getDataType<T>(), sizeof(T), Shape);
  }
 
  const std::string &name() const { return Name; }
  int port() const { return Port; }
  TensorType type() const { return Type; }
  const std::vector<int64_t> &shape() const { return Shape; }
  void setShape(const std::vector<int64_t> &NewShape) {
    Shape = NewShape;
    ElementCount = std::accumulate(Shape.begin(), Shape.end(), 1,
                                   std::multiplies<int64_t>());
  }
  bool operator==(const TensorSpec &Other) const {
    return Name == Other.Name && Port == Other.Port && Type == Other.Type &&
           Shape == Other.Shape;
  }
 
  bool operator!=(const TensorSpec &Other) const { return !(*this == Other); }
 
  size_t getElementCount() const { return ElementCount; }
  size_t getElementByteSize() const { return ElementSize; }
  size_t getTotalTensorBufferSize() const { return ElementCount * ElementSize; }
 
  template <typename T> bool isElementType() const {
    return getDataType<T>() == Type;
  }
 
  TensorSpec(const std::string &NewName, const TensorSpec &Other)
      : TensorSpec(NewName, Other.Port, Other.Type, Other.ElementSize,
                   Other.Shape) {}
 
  void toJSON(llvm::json::OStream &OS) const;
 
private:
  TensorSpec(const std::string &Name, int Port, TensorType Type,
             size_t ElementSize, const std::vector<int64_t> &Shape);
 
  template <typename T> static TensorType getDataType();
 
  std::string Name;
  int Port = 0;
  TensorType Type = TensorType::Invalid;
  std::vector<int64_t> Shape;
  size_t ElementCount = 0;
  size_t ElementSize = 0;
};
 
std::string tensorValueToString(const char *Buffer, const TensorSpec &Spec);
 
llvm::Optional<TensorSpec>
getTensorSpecFromJSON(llvm::LLVMContext &Ctx, const llvm::json::Value &Value);
 
#define TFUTILS_GETDATATYPE_DEF(T, Name)                                       \
  template <> TensorType TensorSpec::getDataType<T>();
SUPPORTED_TENSOR_TYPES(TFUTILS_GETDATATYPE_DEF)
 
#undef TFUTILS_GETDATATYPE_DEF
} // namespace MLBridge
 
#endif