1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20#ifndef TVM_TARGET_DATATYPE_REGISTRY_H_
21#define TVM_TARGET_DATATYPE_REGISTRY_H_
22
23#include <tvm/runtime/packed_func.h>
24#include <tvm/runtime/registry.h>
25
26#include <string>
27#include <unordered_map>
28
29namespace tvm {
30namespace datatype {
31
32/*!
33 * \brief Registry for custom datatypes.
34 *
35 * Adding custom datatypes currently requires two steps:
36 * 1. Register the datatype with the registry via a call to
37 * datatype::Registry::Register. This can also be done in Python
38 * directly---see the TVM globals registered in the corresponding .cc file.
39 * Currently, user should manually choose a type name and a type code,
40 * ensuring that neither conflict with existing types.
41 * 2. Use TVM_REGISTER_GLOBAL to register the lowering functions needed to
42 * lower the custom datatype. In general, these will look like:
43 * For Casts: tvm.datatype.lower.<target>.Cast.<type>.<src_type>
44 * Example: tvm.datatype.lower.llvm.Cast.myfloat.float for a Cast from
45 * float to myfloat.
46 * For intrinsic Calls: tvm.datatype.lower.<target>.Call.intrin.<name>.<type>
47 * Example: tvm.datatype.lower.llvm.Call.intrin.sqrt.myfloat
48 * For other ops: tvm.datatype.lower.<target>.<op>.<type>
49 * Examples: tvm.datatype.lower.llvm.Add.myfloat
50 * tvm.datatype.lower.llvm.FloatImm.posit
51 */
52class Registry {
53 public:
54 /*!
55 * \brief Get the global custom datatype registry singleton
56 */
57 static Registry* Global();
58
59 /*!
60 * \brief Register custom datatype
61 * Register a custom datatype with the given type name and type code. Currently, the type code is
62 * manually allocated by the user, and the user must ensure that no two custom types share the
63 * same code. Generally, this should be straightforward, as the user will be manually registering
64 * all of their custom types.
65 * \param type_name The name of the type, e.g. "posites2"
66 * \param type_code The type code, which should be greater than TVMArgTypeCode::kTVMExtEnd
67 */
68 void Register(const std::string& type_name, uint8_t type_code);
69
70 /*!
71 * \brief Get type code from type name
72 * \param type_name The type name
73 * \return The type code
74 */
75 uint8_t GetTypeCode(const std::string& type_name);
76
77 /*!
78 * \brief Get type name from type code
79 * \param type_code The type code
80 * \return The type name
81 */
82 std::string GetTypeName(uint8_t type_code);
83
84 /*!
85 * \brief Get bool representing whether type is registered, given the type code
86 * \param type_code The type code
87 * \return bool representing whether the type is registered
88 */
89 inline bool GetTypeRegistered(uint8_t type_code) {
90 return code_to_name_.find(type_code) != code_to_name_.end();
91 }
92
93 /*!
94 * \brief Get bool representing whether type is registered, given the type name
95 * \param type_name The type name
96 * \return bool representing whether the type is registered
97 */
98 inline bool GetTypeRegistered(std::string type_name) {
99 return name_to_code_.find(type_name) != name_to_code_.end();
100 }
101
102 private:
103 // TODO(gus) is there a typedef for the code?
104 std::unordered_map<uint8_t, std::string> code_to_name_;
105 std::unordered_map<std::string, uint8_t> name_to_code_;
106};
107
108/*!
109 * \brief Convert scalar value to a custom datatype format
110 * \param type_code The custom datatype to convert to, specified by type code
111 * \param value The floating point value to convert
112 * \return The value, encoded in the bits of a uint64_t
113 */
114uint64_t ConvertConstScalar(uint8_t type_code, double value);
115
116/*!
117 * \brief Get a function returning the minimum value for a datatype.
118 * \param type_code The datatype
119 * \return Function which takes the width of the datatype and returns the min value
120 */
121const runtime::PackedFunc* GetMinFunc(uint8_t type_code);
122
123/*!
124 * \brief Get lowering function for Cast ops
125 * \param target The target we are lowering to, e.g. "llvm"
126 * \param type_code The datatype being cast to
127 * \param src_type_code The datatype being cast from
128 * \return Lowering function for Cast ops for the provided target, type, and source type
129 */
130const runtime::PackedFunc* GetCastLowerFunc(const std::string& target, uint8_t type_code,
131 uint8_t src_type_code);
132
133/*!
134 * \brief Get lowering function for FloatImms
135 * \param target The target we are lowering to, e.g. "llvm"
136 * \param type_code The datatype of the FloatImm
137 * \return Lowering function for FloatImms for the provided target and type
138 */
139const runtime::PackedFunc* GetFloatImmLowerFunc(const std::string& target, uint8_t type_code);
140
141/*!
142 * \brief Get lowering function for intrinsic Calls/pure intrinsic Calls
143 * \param target The target we are lowering to, e.g. "llvm"
144 * \param type_code The datatype of the Call
145 * \param name The intrinsic name
146 * \return Lowering function for intrinsic Calls for the provided target and type
147 */
148const runtime::PackedFunc* GetIntrinLowerFunc(const std::string& target, const std::string& name,
149 uint8_t type_code);
150
151/*!
152 * \brief Get lowering function for other ops
153 * \param target The target we are lowering to, e.g. "llvm"
154 * \param type_code The datatype of the op
155 * \return Lowering function for other ops for the provided target and type
156 */
157#define DEFINE_GET_LOWER_FUNC_(OP) \
158 inline const runtime::PackedFunc* Get##OP##LowerFunc(const std::string& target, \
159 uint8_t type_code) { \
160 return runtime::Registry::Get("tvm.datatype.lower." + target + "." #OP "." + \
161 datatype::Registry::Global()->GetTypeName(type_code)); \
162 }
163
164DEFINE_GET_LOWER_FUNC_(Add)
165DEFINE_GET_LOWER_FUNC_(Sub)
166DEFINE_GET_LOWER_FUNC_(Mul)
167DEFINE_GET_LOWER_FUNC_(Div)
168DEFINE_GET_LOWER_FUNC_(Mod)
169DEFINE_GET_LOWER_FUNC_(Min)
170DEFINE_GET_LOWER_FUNC_(Max)
171DEFINE_GET_LOWER_FUNC_(EQ)
172DEFINE_GET_LOWER_FUNC_(NE)
173DEFINE_GET_LOWER_FUNC_(LT)
174DEFINE_GET_LOWER_FUNC_(LE)
175DEFINE_GET_LOWER_FUNC_(GT)
176DEFINE_GET_LOWER_FUNC_(GE)
177// Later changes may need to add more lowering functions as we support workloads with more ops.
178
179} // namespace datatype
180} // namespace tvm
181
182#endif // TVM_TARGET_DATATYPE_REGISTRY_H_
183