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 | |
29 | namespace tvm { |
30 | namespace 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 | */ |
52 | class 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 | */ |
114 | uint64_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 | */ |
121 | const 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 | */ |
130 | const 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 | */ |
139 | const 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 | */ |
148 | const 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 | |
164 | DEFINE_GET_LOWER_FUNC_(Add) |
165 | DEFINE_GET_LOWER_FUNC_(Sub) |
166 | DEFINE_GET_LOWER_FUNC_(Mul) |
167 | DEFINE_GET_LOWER_FUNC_(Div) |
168 | DEFINE_GET_LOWER_FUNC_(Mod) |
169 | DEFINE_GET_LOWER_FUNC_(Min) |
170 | DEFINE_GET_LOWER_FUNC_(Max) |
171 | DEFINE_GET_LOWER_FUNC_(EQ) |
172 | DEFINE_GET_LOWER_FUNC_(NE) |
173 | DEFINE_GET_LOWER_FUNC_(LT) |
174 | DEFINE_GET_LOWER_FUNC_(LE) |
175 | DEFINE_GET_LOWER_FUNC_(GT) |
176 | DEFINE_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 | |