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
44namespace tvm {
45namespace relay {
46
47TVM_REGISTER_NODE_TYPE(AllocStorageAttrs);
48TVM_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?
53Expr 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
61TVM_REGISTER_GLOBAL("relay.op.memory._make.alloc_storage").set_body_typed(AllocStorage);
62
63bool 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
84RELAY_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
97const Op& MemoryAllocTensorOp() {
98 static const Op& op = Op::Get("memory.alloc_tensor");
99 return op;
100}
101
102Expr 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
117TVM_REGISTER_GLOBAL("relay.op.memory._make.alloc_tensor").set_body_typed(AllocTensor);
118
119std::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
144bool 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
193RELAY_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
207bool 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
215RELAY_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
226static 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
238std::vector<TensorType> FlattenTupleType(const Type& type) {
239 std::vector<TensorType> out;
240 FlattenTupleTypeAux(type, &out);
241 return out;
242}
243
244static 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
256std::vector<Expr> FromTupleType(const Type& type, const Expr& expr) {
257 std::vector<Expr> out;
258 FromTupleTypeAux(type, expr, &out);
259 return out;
260}
261
262static 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.
279Expr 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
290TVM_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
295TVM_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
300TVM_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