LLVM 20.0.0git
TensorSpec.h
Go to the documentation of this file.
1//===- TensorSpec.h - type descriptor for a tensor --------------*- C++ -*-===//
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#ifndef LLVM_ANALYSIS_TENSORSPEC_H
10#define LLVM_ANALYSIS_TENSORSPEC_H
11
12#include "llvm/Config/llvm-config.h"
13
14#include "llvm/ADT/StringMap.h"
15#include "llvm/IR/LLVMContext.h"
16
17#include <memory>
18#include <optional>
19#include <vector>
20
21namespace llvm {
22namespace json {
23class OStream;
24class Value;
25} // namespace json
26
27/// TensorSpec encapsulates the specification of a tensor: its dimensions, or
28/// "shape" (row-major), its type (see TensorSpec::getDataType specializations
29/// for supported types), its name and port (see "TensorFlow: Large-Scale
30/// Machine Learning on Heterogeneous Distributed Systems", section 4.2, para 2:
31/// https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/45166.pdf)
32///
33/// Note that the design is motivated by Tensorflow, but it is not intended to
34/// be Tensorflow-specific.
35///
36/// Known tensor types. The left part is the C type, the
37/// right is a name we can use to identify the type (to implement TensorSpec
38/// equality checks), and to use, if needed, when mapping to an underlying
39/// evaluator's type system. The main requirement is that the C type we use has
40/// the same size and encoding (e.g. endian-ness) as the one used by the
41/// evaluator.
42#define SUPPORTED_TENSOR_TYPES(M) \
43 M(float, Float) \
44 M(double, Double) \
45 M(int8_t, Int8) \
46 M(uint8_t, UInt8) \
47 M(int16_t, Int16) \
48 M(uint16_t, UInt16) \
49 M(int32_t, Int32) \
50 M(uint32_t, UInt32) \
51 M(int64_t, Int64) \
52 M(uint64_t, UInt64)
53
54enum class TensorType {
55 Invalid,
56#define _TENSOR_TYPE_ENUM_MEMBERS(_, Name) Name,
58#undef _TENSOR_TYPE_ENUM_MEMBERS
59 Total
60};
61
62class TensorSpec final {
63public:
64 template <typename T>
65 static TensorSpec createSpec(const std::string &Name,
66 const std::vector<int64_t> &Shape,
67 int Port = 0) {
68 return TensorSpec(Name, Port, getDataType<T>(), sizeof(T), Shape);
69 }
70
71 const std::string &name() const { return Name; }
72 int port() const { return Port; }
73 TensorType type() const { return Type; }
74 const std::vector<int64_t> &shape() const { return Shape; }
75
76 bool operator==(const TensorSpec &Other) const {
77 return Name == Other.Name && Port == Other.Port && Type == Other.Type &&
78 Shape == Other.Shape;
79 }
80
81 bool operator!=(const TensorSpec &Other) const { return !(*this == Other); }
82
83 /// Get the number of elements in a tensor with this shape.
84 size_t getElementCount() const { return ElementCount; }
85 /// Get the size, in bytes, of one element.
86 size_t getElementByteSize() const { return ElementSize; }
87 /// Get the total size of a memory buffer needed to store the whole tensor.
88 size_t getTotalTensorBufferSize() const { return ElementCount * ElementSize; }
89
90 template <typename T> bool isElementType() const {
91 return getDataType<T>() == Type;
92 }
93
94 TensorSpec(const std::string &NewName, const TensorSpec &Other)
95 : TensorSpec(NewName, Other.Port, Other.Type, Other.ElementSize,
96 Other.Shape) {}
97
98 void toJSON(json::OStream &OS) const;
99
100private:
101 TensorSpec(const std::string &Name, int Port, TensorType Type,
102 size_t ElementSize, const std::vector<int64_t> &Shape);
103
104 template <typename T> static TensorType getDataType();
105
106 std::string Name;
107 int Port = 0;
109 std::vector<int64_t> Shape;
110 size_t ElementCount = 0;
111 size_t ElementSize = 0;
112};
113
114/// For debugging.
115std::string tensorValueToString(const char *Buffer, const TensorSpec &Spec);
116
117/// Construct a TensorSpec from a JSON dictionary of the form:
118/// { "name": <string>,
119/// "port": <int>,
120/// "type": <string. Use LLVM's types, e.g. float, double, int64_t>,
121/// "shape": <array of ints> }
122/// For the "type" field, see the C++ primitive types used in
123/// TFUTILS_SUPPORTED_TYPES.
124std::optional<TensorSpec> getTensorSpecFromJSON(LLVMContext &Ctx,
125 const json::Value &Value);
126
127#define TFUTILS_GETDATATYPE_DEF(T, Name) \
128 template <> TensorType TensorSpec::getDataType<T>();
130
131#undef TFUTILS_GETDATATYPE_DEF
132} // namespace llvm
133
134#endif // LLVM_ANALYSIS_TENSORSPEC_H
This file defines the StringMap class.
raw_pwrite_stream & OS
#define TFUTILS_GETDATATYPE_DEF(T, Name)
Definition: TensorSpec.h:127
#define SUPPORTED_TENSOR_TYPES(M)
TensorSpec encapsulates the specification of a tensor: its dimensions, or "shape" (row-major),...
Definition: TensorSpec.h:42
size_t getElementCount() const
Get the number of elements in a tensor with this shape.
Definition: TensorSpec.h:84
const std::string & name() const
Definition: TensorSpec.h:71
TensorType type() const
Definition: TensorSpec.h:73
static TensorSpec createSpec(const std::string &Name, const std::vector< int64_t > &Shape, int Port=0)
Definition: TensorSpec.h:65
const std::vector< int64_t > & shape() const
Definition: TensorSpec.h:74
size_t getElementByteSize() const
Get the size, in bytes, of one element.
Definition: TensorSpec.h:86
bool operator==(const TensorSpec &Other) const
Definition: TensorSpec.h:76
int port() const
Definition: TensorSpec.h:72
bool operator!=(const TensorSpec &Other) const
Definition: TensorSpec.h:81
size_t getTotalTensorBufferSize() const
Get the total size of a memory buffer needed to store the whole tensor.
Definition: TensorSpec.h:88
TensorSpec(const std::string &NewName, const TensorSpec &Other)
Definition: TensorSpec.h:94
void toJSON(json::OStream &OS) const
Definition: TensorSpec.cpp:50
bool isElementType() const
Definition: TensorSpec.h:90
The instances of the Type class are immutable: once they are created, they are never changed.
Definition: Type.h:45
json::OStream allows writing well-formed JSON without materializing all structures as json::Value ahe...
Definition: JSON.h:979
A Value is an JSON value of unknown type.
Definition: JSON.h:288
This is an optimization pass for GlobalISel generic memory operations.
Definition: AddressRanges.h:18
std::optional< TensorSpec > getTensorSpecFromJSON(LLVMContext &Ctx, const json::Value &Value)
Construct a TensorSpec from a JSON dictionary of the form: { "name": <string>, "port": <int>,...
Definition: TensorSpec.cpp:69
@ Other
Any other memory.
std::string tensorValueToString(const char *Buffer, const TensorSpec &Spec)
For debugging.
Definition: TensorSpec.cpp:107
TensorType
Definition: TensorSpec.h:54
_TENSOR_TYPE_ENUM_MEMBERS(_, Name)