1//===- TensorSpec.cpp - tensor type abstraction ---------------------------===//
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// Implementation file for the abstraction of a tensor type, and JSON loading
10// utils.
11//
12//===----------------------------------------------------------------------===//
13#include "llvm/ADT/STLExtras.h"
14#include "llvm/Config/config.h"
15
16#include "llvm/ADT/StringExtras.h"
17#include "llvm/ADT/Twine.h"
18#include "llvm/Analysis/TensorSpec.h"
19#include "llvm/Support/JSON.h"
20#include "llvm/Support/raw_ostream.h"
21#include <array>
22#include <cassert>
23#include <numeric>
24
25using namespace llvm;
26
27namespace llvm {
28
29#define TFUTILS_GETDATATYPE_IMPL(T, E) \
30 template <> TensorType TensorSpec::getDataType<T>() { return TensorType::E; }
31
32SUPPORTED_TENSOR_TYPES(TFUTILS_GETDATATYPE_IMPL)
33
34#undef TFUTILS_GETDATATYPE_IMPL
35
36static std::array<std::string, static_cast<size_t>(TensorType::Total)>
37 TensorTypeNames{"INVALID",
38#define TFUTILS_GETNAME_IMPL(T, _) #T,
39 SUPPORTED_TENSOR_TYPES(TFUTILS_GETNAME_IMPL)
40#undef TFUTILS_GETNAME_IMPL
41 };
42
43StringRef toString(TensorType TT) {
44 return TensorTypeNames[static_cast<size_t>(TT)];
45}
46
47void TensorSpec::toJSON(json::OStream &OS) const {
48 OS.object(Contents: [&]() {
49 OS.attribute(Key: "name", Contents: name());
50 OS.attribute(Key: "type", Contents: toString(TT: type()));
51 OS.attribute(Key: "port", Contents: port());
52 OS.attributeArray(Key: "shape", Contents: [&]() {
53 for (size_t D : shape())
54 OS.value(V: static_cast<int64_t>(D));
55 });
56 });
57}
58
59TensorSpec::TensorSpec(const std::string &Name, int Port, TensorType Type,
60 size_t ElementSize, const std::vector<int64_t> &Shape)
61 : Name(Name), Port(Port), Type(Type), Shape(Shape),
62 ElementCount(std::accumulate(first: Shape.begin(), last: Shape.end(), init: 1,
63 binary_op: std::multiplies<int64_t>())),
64 ElementSize(ElementSize) {}
65
66std::optional<TensorSpec> getTensorSpecFromJSON(LLVMContext &Ctx,
67 const json::Value &Value) {
68 auto EmitError =
69 [&](const llvm::Twine &Message) -> std::optional<TensorSpec> {
70 std::string S;
71 llvm::raw_string_ostream OS(S);
72 OS << Value;
73 Ctx.emitError(ErrorStr: "Unable to parse JSON Value as spec (" + Message + "): " + S);
74 return std::nullopt;
75 };
76 // FIXME: accept a Path as a parameter, and use it for error reporting.
77 json::Path::Root Root("tensor_spec");
78 json::ObjectMapper Mapper(Value, Root);
79 if (!Mapper)
80 return EmitError("Value is not a dict");
81
82 std::string TensorName;
83 int TensorPort = -1;
84 std::string TensorType;
85 std::vector<int64_t> TensorShape;
86
87 if (!Mapper.map<std::string>(Prop: "name", Out&: TensorName))
88 return EmitError("'name' property not present or not a string");
89 if (!Mapper.map<std::string>(Prop: "type", Out&: TensorType))
90 return EmitError("'type' property not present or not a string");
91 if (!Mapper.map<int>(Prop: "port", Out&: TensorPort))
92 return EmitError("'port' property not present or not an int");
93 if (!Mapper.map<std::vector<int64_t>>(Prop: "shape", Out&: TensorShape))
94 return EmitError("'shape' property not present or not an int array");
95
96#define PARSE_TYPE(T, E) \
97 if (TensorType == #T) \
98 return TensorSpec::createSpec<T>(TensorName, TensorShape, TensorPort);
99 SUPPORTED_TENSOR_TYPES(PARSE_TYPE)
100#undef PARSE_TYPE
101 return std::nullopt;
102}
103
104std::string tensorValueToString(const char *Buffer, const TensorSpec &Spec) {
105 switch (Spec.type()) {
106#define _IMR_DBG_PRINTER(T, N) \
107 case TensorType::N: { \
108 const T *TypedBuff = reinterpret_cast<const T *>(Buffer); \
109 auto R = llvm::make_range(TypedBuff, TypedBuff + Spec.getElementCount()); \
110 return llvm::join( \
111 llvm::map_range(R, [](T V) { return std::to_string(V); }), ","); \
112 }
113 SUPPORTED_TENSOR_TYPES(_IMR_DBG_PRINTER)
114#undef _IMR_DBG_PRINTER
115 case TensorType::Total:
116 case TensorType::Invalid:
117 llvm_unreachable("invalid tensor type");
118 }
119 // To appease warnings about not all control paths returning a value.
120 return "";
121}
122
123} // namespace llvm
124