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
|
#pragma once
#ifdef __GNUC__
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-parameter"
#endif
//===- TensorSpec.h - type descriptor for a tensor --------------*- 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_TENSORSPEC_H
#define LLVM_ANALYSIS_TENSORSPEC_H
#include "llvm/Config/llvm-config.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/Support/JSON.h"
#include <memory>
#include <optional>
#include <vector>
namespace llvm {
/// 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)
///
/// Known tensor types. The left part is the C type, the right is a name we
/// can use to identify the type (to implement TensorSpec equality checks), and
/// to use, if needed, when mapping to an underlying evaluator's type system.
/// The main requirement is that the C type we use has the same size and
/// encoding (e.g. endian-ness) as the one used by the evaluator.
#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 = 0) {
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; }
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); }
/// 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 { return ElementSize; }
/// Get the total size of a memory buffer needed to store the whole tensor.
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(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;
};
/// 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.
std::optional<TensorSpec> getTensorSpecFromJSON(LLVMContext &Ctx,
const 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 llvm
#endif // LLVM_ANALYSIS_TENSORSPEC_H
#ifdef __GNUC__
#pragma GCC diagnostic pop
#endif
|
