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 | /*! |
21 | * \file src/relay/op/memory/memory.cc |
22 | * \brief Operators for manifest shape-aware memory allocation in Relay. |
23 | */ |
24 | |
25 | #include "memory.h" |
26 | |
27 | #include <tvm/node/node.h> |
28 | #include <tvm/relay/attrs/memory.h> |
29 | #include <tvm/relay/expr.h> |
30 | #include <tvm/relay/op.h> |
31 | #include <tvm/relay/op_attr_types.h> |
32 | #include <tvm/runtime/data_type.h> |
33 | #include <tvm/topi/elemwise.h> |
34 | |
35 | #include <utility> |
36 | #include <vector> |
37 | |
38 | #include "../../transforms/infer_layout_utils.h" |
39 | #include "../annotation/annotation.h" |
40 | #include "../op_common.h" |
41 | #include "../type_relations.h" |
42 | #include "on_device.h" |
43 | |
44 | namespace tvm { |
45 | namespace relay { |
46 | |
47 | TVM_REGISTER_NODE_TYPE(AllocStorageAttrs); |
48 | TVM_REGISTER_NODE_TYPE(AllocTensorAttrs); |
49 | |
50 | // The passing value in attrs and args doesn't seem super great. |
51 | // We should consider a better solution, i.e the type relation |
52 | // being able to see the arguments as well? |
53 | Expr AllocStorage(Expr size, Expr alignment, VirtualDevice virtual_device, DataType dtype_hint) { |
54 | auto attrs = make_object<AllocStorageAttrs>(); |
55 | attrs->dtype = dtype_hint; |
56 | attrs->virtual_device = std::move(virtual_device); |
57 | static const Op& op = Op::Get("memory.alloc_storage" ); |
58 | return Call(op, {std::move(size), std::move(alignment)}, Attrs(std::move(attrs)), {}); |
59 | } |
60 | |
61 | TVM_REGISTER_GLOBAL("relay.op.memory._make.alloc_storage" ).set_body_typed(AllocStorage); |
62 | |
63 | bool AllocStorageRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, |
64 | const TypeReporter& reporter) { |
65 | ICHECK_EQ(types.size(), 3u); |
66 | auto size_type = types[0]; |
67 | auto tensor_type = size_type.as<TensorTypeNode>(); |
68 | ICHECK(tensor_type != nullptr); |
69 | ICHECK_EQ(tensor_type->dtype, DataType::Int(64)); |
70 | ICHECK_EQ(tensor_type->shape.size(), 0); |
71 | auto align_type = types[1]; |
72 | auto align_ttype = align_type.as<TensorTypeNode>(); |
73 | ICHECK(align_ttype != nullptr); |
74 | ICHECK_EQ(align_ttype->dtype, DataType::Int(64)); |
75 | ICHECK_EQ(align_ttype->shape.size(), 0); |
76 | auto mod = reporter->GetModule(); |
77 | ICHECK(mod.defined()); |
78 | auto storage_name = mod->GetGlobalTypeVar("Storage" ); |
79 | auto storage = TypeCall(storage_name, {}); |
80 | reporter->Assign(types[2], storage); |
81 | return true; |
82 | } |
83 | |
84 | RELAY_REGISTER_OP("memory.alloc_storage" ) |
85 | .describe(R"code(Explicitly allocate storage to be used by tensors.)code" TVM_ADD_FILELINE) |
86 | .set_num_inputs(2) |
87 | .add_argument("size" , "Tensor" , "The size of the storage to allocate." ) |
88 | .add_argument("alignment" , "Tensor" , "The alignment of the storage." ) |
89 | .add_type_rel("AllocStorage" , AllocStorageRel) |
90 | .set_attrs_type_key("relay.attrs.AllocStorageAttrs" ) |
91 | .set_support_level(10) |
92 | .set_attr<TOpPattern>("TOpPattern" , kOpaque) |
93 | .set_attr<TOpIsStateful>("TOpIsStateful" , false) |
94 | .set_attr<TNonComputational>("TNonComputational" , true) |
95 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , ElemwiseArbitraryLayout); |
96 | |
97 | const Op& MemoryAllocTensorOp() { |
98 | static const Op& op = Op::Get("memory.alloc_tensor" ); |
99 | return op; |
100 | } |
101 | |
102 | Expr AllocTensor(Expr storage, Expr offset, Expr shape, DataType dtype, |
103 | Array<IndexExpr> assert_shape) { |
104 | auto attrs = make_object<AllocTensorAttrs>(); |
105 | attrs->dtype = dtype; |
106 | if (assert_shape.defined()) { |
107 | attrs->assert_shape = assert_shape; |
108 | } else { |
109 | // Look through any on_device for the shape argument expression. |
110 | const auto* constant_node = AsIgnoringOnDevice<ConstantNode>(shape); |
111 | ICHECK(constant_node); |
112 | attrs->const_shape = GetRef<Constant>(constant_node); |
113 | } |
114 | return Call(MemoryAllocTensorOp(), {storage, offset, shape}, Attrs(attrs), {}); |
115 | } |
116 | |
117 | TVM_REGISTER_GLOBAL("relay.op.memory._make.alloc_tensor" ).set_body_typed(AllocTensor); |
118 | |
119 | std::vector<int64_t> FromConstShape(Constant konst) { |
120 | runtime::NDArray shape = konst->data; |
121 | std::vector<int64_t> raw_shape; |
122 | ICHECK_EQ(shape->ndim, 1u); |
123 | ICHECK_EQ(shape->dtype.code, 0U) << "The dtype of constant shape must be int32 or int64, but got " |
124 | << runtime::DLDataType2String(shape->dtype); |
125 | ICHECK(shape->dtype.bits == 64 || shape->dtype.bits == 32) |
126 | << "The dtype of constant shape must be int32 or int64, but got" |
127 | << runtime::DLDataType2String(shape->dtype); |
128 | |
129 | if (shape->dtype.bits == 32) { |
130 | const int32_t* int_ptr = reinterpret_cast<int32_t*>(shape->data); |
131 | for (auto i = 0; i < shape->shape[0]; i++) { |
132 | raw_shape.push_back(int_ptr[i]); |
133 | } |
134 | } else if (shape->dtype.bits == 64) { |
135 | const int64_t* int_ptr = reinterpret_cast<int64_t*>(shape->data); |
136 | for (auto i = 0; i < shape->shape[0]; i++) { |
137 | raw_shape.push_back(int_ptr[i]); |
138 | } |
139 | } |
140 | |
141 | return raw_shape; |
142 | } |
143 | |
144 | bool AllocTensorRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, |
145 | const TypeReporter& reporter) { |
146 | ICHECK_EQ(types.size(), 4u); |
147 | auto alloc_attrs = attrs.as<AllocTensorAttrs>(); |
148 | ICHECK(alloc_attrs != nullptr) << "must be alloc_tensor attributes" ; |
149 | // First argument should be storage. |
150 | auto mod = reporter->GetModule(); |
151 | ICHECK(mod.defined()); |
152 | auto storage_name = mod->GetGlobalTypeVar("Storage" ); |
153 | auto storage = relay::TypeCall(storage_name, {}); |
154 | reporter->Assign(types[0], storage); |
155 | // Second argument should be the offset. |
156 | auto offset_type = types[1].as<TensorTypeNode>(); |
157 | ICHECK(offset_type != nullptr) << "must be a scalar type" ; |
158 | |
159 | // Third argument should be shape tensor. |
160 | auto tt = types[2].as<TensorTypeNode>(); |
161 | ICHECK(tt != nullptr) << "must be tensor type" ; |
162 | |
163 | // Be careful about having to allocate scalars. |
164 | int64_t dims = 0; |
165 | if (tt->shape.size() != 0) { |
166 | auto rank = tt->shape[0].as<tvm::IntImmNode>(); |
167 | ICHECK(rank != nullptr); |
168 | dims = rank->value; |
169 | } |
170 | |
171 | // Constant node case. |
172 | Type alloc_type; |
173 | if (alloc_attrs->const_shape.defined()) { |
174 | auto con = alloc_attrs->const_shape; |
175 | auto sh = FromConstShape(con); |
176 | ICHECK_EQ(sh.size(), dims); |
177 | Array<IndexExpr> out_shape; |
178 | for (auto i = 0u; i < dims; i++) { |
179 | out_shape.push_back(tvm::Integer(sh[i])); |
180 | } |
181 | alloc_type = TensorType(out_shape, alloc_attrs->dtype); |
182 | } else { |
183 | ICHECK(alloc_attrs->assert_shape.defined()) |
184 | << "the assert_shape must be set when const_shape is not" ; |
185 | alloc_type = TensorType(alloc_attrs->assert_shape, alloc_attrs->dtype); |
186 | return true; |
187 | } |
188 | |
189 | reporter->Assign(types[3], alloc_type); |
190 | return true; |
191 | } |
192 | |
193 | RELAY_REGISTER_OP("memory.alloc_tensor" ) |
194 | .describe(R"code(Explicitly allocate storage to be used by tensors.)code" TVM_ADD_FILELINE) |
195 | .set_num_inputs(3) |
196 | .add_argument("storage" , "Storage" , "The storage to allocate from." ) |
197 | .add_argument("offset" , "Tensor" , "The offset into the backing storage." ) |
198 | .add_argument("shape" , "Tensor" , "The shape of the tensor to allocate." ) |
199 | .add_type_rel("AllocTensor" , AllocTensorRel) |
200 | .set_attrs_type_key("relay.attrs.AllocTensorAttrs" ) |
201 | .set_support_level(10) |
202 | .set_attr<TOpPattern>("TOpPattern" , kOpaque) |
203 | .set_attr<TOpIsStateful>("TOpIsStateful" , false) |
204 | .set_attr<TNonComputational>("TNonComputational" , true) |
205 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , ElemwiseArbitraryLayout); |
206 | |
207 | bool KillRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, |
208 | const TypeReporter& reporter) { |
209 | ICHECK_EQ(types.size(), 2u); |
210 | // TODO(@jroesch): should only support tensors. |
211 | reporter->Assign(types[1], TupleType::Empty()); |
212 | return true; |
213 | } |
214 | |
215 | RELAY_REGISTER_OP("memory.kill" ) |
216 | .describe(R"code(Mark a variable for release to the allocator.)code" TVM_ADD_FILELINE) |
217 | .set_num_inputs(1) |
218 | .add_argument("to_free" , "Variable" , "The variable to free." ) |
219 | .add_type_rel("Kill" , KillRel) |
220 | .set_support_level(10) |
221 | .set_attr<TOpPattern>("TOpPattern" , kOpaque) |
222 | .set_attr<TOpIsStateful>("TOpIsStateful" , true) |
223 | .set_attr<TNonComputational>("TNonComputational" , true) |
224 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , ElemwiseArbitraryLayout); |
225 | |
226 | static void FlattenTupleTypeAux(const Type& type, std::vector<TensorType>* out) { |
227 | if (auto tt = type.as<TensorTypeNode>()) { |
228 | out->push_back(GetRef<TensorType>(tt)); |
229 | } else if (auto tuple_ty = type.as<TupleTypeNode>()) { |
230 | for (auto field : tuple_ty->fields) { |
231 | FlattenTupleTypeAux(field, out); |
232 | } |
233 | } else { |
234 | LOG(FATAL) << "unsupported " << type; |
235 | } |
236 | } |
237 | |
238 | std::vector<TensorType> FlattenTupleType(const Type& type) { |
239 | std::vector<TensorType> out; |
240 | FlattenTupleTypeAux(type, &out); |
241 | return out; |
242 | } |
243 | |
244 | static void FromTupleTypeAux(const Type& type, const Expr& expr, std::vector<Expr>* out) { |
245 | if (type.as<TensorTypeNode>()) { |
246 | out->push_back(expr); |
247 | } else if (auto tuple_ty = type.as<TupleTypeNode>()) { |
248 | for (size_t i = 0; i < tuple_ty->fields.size(); i++) { |
249 | FromTupleTypeAux(tuple_ty->fields[i], TupleGetItem(expr, i), out); |
250 | } |
251 | } else { |
252 | LOG(FATAL) << "unsupported " << type; |
253 | } |
254 | } |
255 | |
256 | std::vector<Expr> FromTupleType(const Type& type, const Expr& expr) { |
257 | std::vector<Expr> out; |
258 | FromTupleTypeAux(type, expr, &out); |
259 | return out; |
260 | } |
261 | |
262 | static void ToTupleTypeAux(const Type& type, const std::vector<Expr>& exprs, int* index, |
263 | std::vector<Expr>* out) { |
264 | if (type.as<TensorTypeNode>()) { |
265 | out->push_back(exprs[*index]); |
266 | *index += 1; |
267 | } else if (auto tuple_ty = type.as<TupleTypeNode>()) { |
268 | std::vector<Expr> tuple_out; |
269 | for (size_t i = 0; i < tuple_ty->fields.size(); i++) { |
270 | ToTupleTypeAux(tuple_ty->fields[i], exprs, index, &tuple_out); |
271 | } |
272 | out->push_back(Tuple(tuple_out)); |
273 | } else { |
274 | LOG(FATAL) << "unsupported " << type; |
275 | } |
276 | } |
277 | |
278 | // Pack the sequence of expressions according to the provided TupleType. |
279 | Expr ToTupleType(const Type& t, const std::vector<Expr>& exprs) { |
280 | if (t.as<TensorTypeNode>() && exprs.size() == 1) { |
281 | return exprs[0]; |
282 | } else { |
283 | std::vector<Expr> out; |
284 | int index = 0; |
285 | ToTupleTypeAux(t, exprs, &index, &out); |
286 | return out[0]; |
287 | } |
288 | } |
289 | |
290 | TVM_REGISTER_GLOBAL("relay.op.memory._make.FlattenTupleType" ).set_body_typed([](Type type) { |
291 | auto types = FlattenTupleType(type); |
292 | return Array<Type>(types.begin(), types.end()); |
293 | }); |
294 | |
295 | TVM_REGISTER_GLOBAL("relay.op.memory._make.FromTupleType" ).set_body_typed([](Type type, Expr expr) { |
296 | auto exprs = FromTupleType(type, expr); |
297 | return Array<Expr>(exprs.begin(), exprs.end()); |
298 | }); |
299 | |
300 | TVM_REGISTER_GLOBAL("relay.op.memory._make.ToTupleType" ) |
301 | .set_body_typed([](Type t, Array<Expr> array) { |
302 | return ToTupleType(t, std::vector<Expr>(array.begin(), array.end())); |
303 | }); |
304 | |
305 | } // namespace relay |
306 | } // namespace tvm |
307 | |