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 pooling.cc |
22 | * \brief Pooling operators |
23 | */ |
24 | #include "pooling.h" |
25 | |
26 | #include <tvm/relay/attrs/nn.h> |
27 | #include <tvm/relay/op.h> |
28 | #include <tvm/relay/op_attr_types.h> |
29 | #include <tvm/tir/data_layout.h> |
30 | #include <tvm/topi/nn/pooling.h> |
31 | |
32 | #include <vector> |
33 | |
34 | #include "../../transforms/infer_layout_utils.h" |
35 | |
36 | namespace tvm { |
37 | namespace relay { |
38 | |
39 | // relay.nn.max_pool2d & relay.nn.avg_pool2d |
40 | TVM_REGISTER_NODE_TYPE(MaxPool2DAttrs); |
41 | TVM_REGISTER_NODE_TYPE(AvgPool2DAttrs); |
42 | |
43 | template <typename T> |
44 | InferCorrectLayoutOutput PoolInferCorrectLayout(const Attrs& attrs, |
45 | const Array<Layout>& new_in_layouts, |
46 | const Array<Layout>& old_in_layouts, |
47 | const Array<tvm::relay::Type>& old_in_types) { |
48 | const auto* attrs_ptr = attrs.as<T>(); |
49 | ICHECK(attrs_ptr); |
50 | ObjectPtr<T> params = make_object<T>(*attrs_ptr); |
51 | |
52 | if (params->out_layout != "" ) { |
53 | // when users specify the out_layout of pooling, follow user's preference |
54 | ICHECK_EQ(params->layout, params->out_layout) |
55 | << "Pooling input/output layouts mismatch: " << params->layout << " vs. " |
56 | << params->out_layout; |
57 | } else if (new_in_layouts.defined()) { |
58 | // the pooling is using an inferred layout (i.e., new_in_layouts[0]) given by relay caller |
59 | ICHECK_EQ(new_in_layouts.size(), 1); |
60 | params->layout = new_in_layouts[0].name(); |
61 | } |
62 | |
63 | return InferCorrectLayoutOutput({params->layout}, {params->layout}, Attrs(params)); |
64 | } |
65 | |
66 | IndexExpr calculate_pool_dimension(IndexExpr in_dimension, IndexExpr pad_amount, |
67 | IndexExpr pool_size, IndexExpr dilation, IndexExpr stride_size, |
68 | bool ceil_mode) { |
69 | IndexExpr numerator = in_dimension + pad_amount - ((pool_size - 1) * dilation + 1); |
70 | IndexExpr denominator = stride_size; |
71 | |
72 | // Emulate the behavior of running ceil on numerator / denominator rather than floor |
73 | if (ceil_mode) { |
74 | numerator += denominator - 1; |
75 | } |
76 | |
77 | return numerator / denominator + 1; |
78 | } |
79 | |
80 | template <typename AttrType> |
81 | bool Pool2DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, |
82 | const TypeReporter& reporter) { |
83 | ICHECK_EQ(types.size(), 2); |
84 | const auto* data = types[0].as<TensorTypeNode>(); |
85 | |
86 | if (data == nullptr) return false; |
87 | |
88 | const auto dshape = data->shape; |
89 | ICHECK_GE(dshape.size(), 2U) |
90 | << "Pool2D only support input >= 2-D: input must have height and width" ; |
91 | const auto param = attrs.as<AttrType>(); |
92 | ICHECK(param != nullptr); |
93 | |
94 | Layout layout(param->layout); |
95 | ICHECK(layout.Contains(LayoutAxis::Get('H')) && layout.Contains(LayoutAxis::Get('W')) && |
96 | !layout.Contains(LayoutAxis::Get('h')) && !layout.Contains(LayoutAxis::Get('w'))) |
97 | << "Invalid layout " << layout << ". Pool2D layout must have H and W, which cannot be split" ; |
98 | |
99 | const auto hidx = layout.IndexOf(LayoutAxis::Get('H')); |
100 | const auto widx = layout.IndexOf(LayoutAxis::Get('W')); |
101 | |
102 | IndexExpr pad_h, pad_w; |
103 | if (param->padding.size() == 1) { |
104 | pad_h = param->padding[0] * 2; |
105 | pad_w = param->padding[0] * 2; |
106 | } else if (param->padding.size() == 2) { |
107 | // (top, left) |
108 | pad_h = param->padding[0] * 2; |
109 | pad_w = param->padding[1] * 2; |
110 | } else if (param->padding.size() == 4) { |
111 | // (top, left, bottom, right) |
112 | pad_h = param->padding[0] + param->padding[2]; |
113 | pad_w = param->padding[1] + param->padding[3]; |
114 | } else { |
115 | return false; |
116 | } |
117 | |
118 | std::vector<IndexExpr> oshape(dshape.begin(), dshape.end()); |
119 | |
120 | if (dshape[hidx].as<tir::AnyNode>()) { |
121 | oshape[hidx] = dshape[hidx]; |
122 | } else { |
123 | oshape[hidx] = |
124 | calculate_pool_dimension(dshape[hidx], pad_h, param->pool_size[0], param->dilation[0], |
125 | param->strides[0], param->ceil_mode); |
126 | } |
127 | if (dshape[widx].as<tir::AnyNode>()) { |
128 | oshape[widx] = dshape[widx]; |
129 | } else { |
130 | oshape[widx] = |
131 | calculate_pool_dimension(dshape[widx], pad_w, param->pool_size[1], param->dilation[1], |
132 | param->strides[1], param->ceil_mode); |
133 | } |
134 | |
135 | // assign output type |
136 | reporter->Assign(types[1], TensorType(oshape, data->dtype)); |
137 | return true; |
138 | } |
139 | |
140 | template <typename AttrType, topi::nn::PoolType mode> |
141 | Array<te::Tensor> Pool2DCompute(const Attrs& attrs, const Array<te::Tensor>& inputs, |
142 | const Type& out_type) { |
143 | static const Layout kNCHW("NCHW" ); |
144 | const auto* param = attrs.as<AttrType>(); |
145 | ICHECK(param != nullptr); |
146 | auto pool_size = param->pool_size; |
147 | auto strides = param->strides; |
148 | auto dilation = param->dilation; |
149 | auto padding = param->padding; |
150 | auto ceil_mode = param->ceil_mode; |
151 | Layout layout(param->layout); |
152 | Layout out_layout(param->out_layout); |
153 | |
154 | ICHECK(tir::BijectiveLayout(layout, kNCHW).defined()) |
155 | << "max_pool2d currently only supports layouts that are convertible from NCHW" ; |
156 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('h')), -1) |
157 | << "max_pool2d does not support input split on height" ; |
158 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('w')), -1) |
159 | << "max_pool2d does not support input split on width" ; |
160 | |
161 | ICHECK(inputs[0].ndim() == 4U || inputs[0].ndim() == 5U || inputs[0].ndim() == 6U) |
162 | << "Pool2D only support 4-D input (e.g., NCHW)" |
163 | << " or 5-D input (e.g. NCHWc on for vector instructions)" |
164 | << " or 6-D input (e.g. NCHWnc for tensor accelerators)" ; |
165 | |
166 | if (param->padding.size() == 1) { |
167 | padding.push_back(padding[0]); |
168 | padding.push_back(padding[0]); |
169 | padding.push_back(padding[0]); |
170 | } else if (param->padding.size() == 2) { |
171 | padding.push_back(padding[0]); |
172 | padding.push_back(padding[1]); |
173 | } |
174 | if (mode == topi::nn::kAvgPool) { |
175 | bool count_include_pad = reinterpret_cast<const AvgPool2DAttrs*>(param)->count_include_pad; |
176 | return Array<te::Tensor>{topi::nn::pool2d(inputs[0], pool_size, strides, dilation, padding, |
177 | mode, ceil_mode, layout.name(), count_include_pad)}; |
178 | } else { |
179 | return Array<te::Tensor>{topi::nn::pool2d(inputs[0], pool_size, strides, dilation, padding, |
180 | mode, ceil_mode, layout.name())}; |
181 | } |
182 | } |
183 | |
184 | TVM_REGISTER_GLOBAL("relay.op.nn._make.max_pool2d" ) |
185 | .set_body_typed([](Expr data, Array<IndexExpr> pool_size, Array<IndexExpr> strides, |
186 | Array<IndexExpr> dilation, Array<IndexExpr> padding, String layout, |
187 | String out_layout, bool ceil_mode) { |
188 | return MakeMaxPool<MaxPool2DAttrs>(data, pool_size, strides, dilation, padding, layout, |
189 | out_layout, ceil_mode, "nn.max_pool2d" ); |
190 | }); |
191 | |
192 | RELAY_REGISTER_OP("nn.max_pool2d" ) |
193 | .describe(R"code(Max pooling operation for two dimensional data. |
194 | |
195 | - **data**: This depends on the `layout` parameter. Input is 4D array of shape |
196 | (batch_size, channels, height, width) if `layout` is `NCHW`. |
197 | - **out**: This depends on the `layout` parameter. Output is 4D array of shape |
198 | (batch_size, channels, out_height, out_width) if `layout` is `NCHW`. |
199 | out_height and out_width are calculated as:: |
200 | |
201 | out_height = floor((height+padding[0]+padding[2]-pool_size[0])/strides[0])+1 |
202 | out_width = floor((width+padding[1]+padding[3]-pool_size[1])/strides[1])+1 |
203 | |
204 | where padding will be an expanded array based on number of values passed as:: |
205 | one int : all sides same padding used. |
206 | two int : bottom, right use same as top and left. |
207 | four int: padding width in the order of (top, left, bottom, right). |
208 | |
209 | When `ceil_mode` is `True`, ceil will be used instead of floor in this |
210 | equation. |
211 | |
212 | )code" TVM_ADD_FILELINE) |
213 | .set_attrs_type<MaxPool2DAttrs>() |
214 | .set_num_inputs(1) |
215 | .add_argument("data" , "Tensor" , "The input tensor." ) |
216 | .set_support_level(2) |
217 | .add_type_rel("MaxPool2D" , Pool2DRel<MaxPool2DAttrs>) |
218 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , PoolInferCorrectLayout<MaxPool2DAttrs>) |
219 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
220 | .set_attr<FTVMCompute>("FTVMCompute" , Pool2DCompute<MaxPool2DAttrs, topi::nn::kMaxPool>); |
221 | |
222 | // AvgPool2D |
223 | TVM_REGISTER_GLOBAL("relay.op.nn._make.avg_pool2d" ) |
224 | .set_body_typed([](Expr data, Array<IndexExpr> pool_size, Array<IndexExpr> strides, |
225 | Array<IndexExpr> dilation, Array<IndexExpr> padding, String layout, |
226 | String out_layout, bool ceil_mode, bool count_include_pad) { |
227 | return MakeAvgPool<AvgPool2DAttrs>(data, pool_size, strides, dilation, padding, layout, |
228 | out_layout, ceil_mode, count_include_pad, "nn.avg_pool2d" ); |
229 | }); |
230 | |
231 | RELAY_REGISTER_OP("nn.avg_pool2d" ) |
232 | .describe(R"code( |
233 | Average pooling operation for one dimensional data. |
234 | |
235 | - **data**: This depends on the `layout` parameter. Input is 4D array of shape |
236 | (batch_size, channels, height, width) if `layout` is `NCHW`. |
237 | - **out**: This depends on the `layout` parameter. Output is 4D array of shape |
238 | (batch_size, channels, out_height, out_width) if `layout` is `NCHW`. |
239 | out_height and out_width are calculated as:: |
240 | |
241 | out_height = floor((height+padding[0]+padding[2]-pool_size[0])/strides[0])+1 |
242 | out_width = floor((width+padding[1]+padding[3]-pool_size[1])/strides[1])+1 |
243 | |
244 | where padding will be an expanded array based on number of values passed as:: |
245 | one int : all sides same padding used. |
246 | two int : bottom, right use same as top and left. |
247 | four int: padding width in the order of (top, left, bottom, right). |
248 | |
249 | When `ceil_mode` is `True`, ceil will be used instead of floor in this |
250 | equation. |
251 | |
252 | )code" TVM_ADD_FILELINE) |
253 | .set_attrs_type<AvgPool2DAttrs>() |
254 | .set_num_inputs(1) |
255 | .add_argument("data" , "Tensor" , "The input tensor." ) |
256 | .set_support_level(2) |
257 | .add_type_rel("AvgPool2D" , Pool2DRel<AvgPool2DAttrs>) |
258 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , PoolInferCorrectLayout<AvgPool2DAttrs>) |
259 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
260 | .set_attr<FTVMCompute>("FTVMCompute" , Pool2DCompute<AvgPool2DAttrs, topi::nn::kAvgPool>); |
261 | |
262 | // relay.nn.global_pool_2d & relay.nn.max_pool_2d |
263 | TVM_REGISTER_NODE_TYPE(GlobalPool2DAttrs); |
264 | |
265 | bool GlobalPool2DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, |
266 | const TypeReporter& reporter) { |
267 | ICHECK_EQ(types.size(), 2); |
268 | const auto* data = types[0].as<TensorTypeNode>(); |
269 | if (data == nullptr) { |
270 | return false; |
271 | } |
272 | const auto dshape = data->shape; |
273 | ICHECK_GE(dshape.size(), 2U) |
274 | << "Pool2D only support input >= 2-D: input must have height and width" ; |
275 | const auto param = attrs.as<GlobalPool2DAttrs>(); |
276 | ICHECK(param != nullptr); |
277 | |
278 | Layout layout(param->layout); |
279 | ICHECK(layout.Contains(LayoutAxis::Get('H')) && layout.Contains(LayoutAxis::Get('W')) && |
280 | !layout.Contains(LayoutAxis::Get('h')) && !layout.Contains(LayoutAxis::Get('w'))) |
281 | << "Invalid layout " << layout << ". Pool2D layout must have H and W, which cannot be split" ; |
282 | |
283 | const auto hidx = layout.IndexOf(LayoutAxis::Get('H')); |
284 | const auto widx = layout.IndexOf(LayoutAxis::Get('W')); |
285 | Array<IndexExpr> oshape(dshape); |
286 | oshape.Set(hidx, 1); |
287 | oshape.Set(widx, 1); |
288 | |
289 | // assign output type |
290 | reporter->Assign(types[1], TensorType(oshape, data->dtype)); |
291 | return true; |
292 | } |
293 | |
294 | template <topi::nn::PoolType mode> |
295 | Array<te::Tensor> GlobalPool2DCompute(const Attrs& attrs, const Array<te::Tensor>& inputs, |
296 | const Type& out_type) { |
297 | static const Layout kNCHW("NCHW" ); |
298 | const auto* param = attrs.as<GlobalPool2DAttrs>(); |
299 | ICHECK(param != nullptr); |
300 | Layout layout(param->layout); |
301 | ICHECK(tir::BijectiveLayout(layout, kNCHW).defined()) |
302 | << "global_avg_pool2d currently only supports layouts that are convertible from NCHW" ; |
303 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('h')), -1) |
304 | << "global_avg_pool2d does not support input split on height" ; |
305 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('w')), -1) |
306 | << "global_avg_pool2d does not support input split on width" ; |
307 | |
308 | ICHECK(inputs[0].ndim() == 4U || inputs[0].ndim() == 5U) |
309 | << "Pool2D only support 4-D input (e.g., NCHW)" |
310 | << " or 5-D input (last dimension is a split of channel)" ; |
311 | return Array<te::Tensor>{topi::nn::global_pool(inputs[0], mode, layout.name())}; |
312 | } |
313 | |
314 | Expr MakeGlobalAvgPool2D(Expr data, String layout, String out_layout) { |
315 | auto attrs = make_object<GlobalPool2DAttrs>(); |
316 | attrs->layout = std::move(layout); |
317 | attrs->out_layout = std::move(out_layout); |
318 | static const Op& op = Op::Get("nn.global_avg_pool2d" ); |
319 | return Call(op, {data}, Attrs(attrs), {}); |
320 | } |
321 | |
322 | TVM_REGISTER_GLOBAL("relay.op.nn._make.global_avg_pool2d" ).set_body_typed(MakeGlobalAvgPool2D); |
323 | |
324 | // GlobalAvgPool |
325 | RELAY_REGISTER_OP("nn.global_avg_pool2d" ) |
326 | .describe(R"code(Global average pooling operation for 2D data. |
327 | |
328 | - **data**: This depends on the `layout` parameter. Input is 4D array of shape |
329 | (batch_size, channels, height, width) if `layout` is `NCHW`. |
330 | - **out**: This depends on the `layout` parameter. Output is 4D array of shape |
331 | (batch_size, channels, 1, 1) if `layout` is `NCHW`. |
332 | |
333 | )code" TVM_ADD_FILELINE) |
334 | .set_attrs_type<GlobalPool2DAttrs>() |
335 | .set_num_inputs(1) |
336 | .add_argument("data" , "Tensor" , "The input tensor." ) |
337 | .set_support_level(2) |
338 | .add_type_rel("GlobalAvgPool2D" , GlobalPool2DRel) |
339 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , PoolInferCorrectLayout<GlobalPool2DAttrs>) |
340 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
341 | .set_attr<FTVMCompute>("FTVMCompute" , GlobalPool2DCompute<topi::nn::kAvgPool>); |
342 | |
343 | // GlobalMaxPool |
344 | Expr MakeGlobalMaxPool2D(Expr data, String layout, String out_layout) { |
345 | auto attrs = make_object<GlobalPool2DAttrs>(); |
346 | attrs->layout = std::move(layout); |
347 | attrs->out_layout = std::move(out_layout); |
348 | static const Op& op = Op::Get("nn.global_max_pool2d" ); |
349 | return Call(op, {data}, Attrs(attrs), {}); |
350 | } |
351 | |
352 | TVM_REGISTER_GLOBAL("relay.op.nn._make.global_max_pool2d" ).set_body_typed(MakeGlobalMaxPool2D); |
353 | |
354 | RELAY_REGISTER_OP("nn.global_max_pool2d" ) |
355 | .describe(R"code(Global max pooling operation for 2D data. |
356 | |
357 | - **data**: This depends on the `layout` parameter. Input is 4D array of shape |
358 | (batch_size, channels, height, width) if `layout` is `NCHW`. |
359 | - **out**: This depends on the `layout` parameter. Output is 4D array of shape |
360 | (batch_size, channels, 1, 1) if `layout` is `NCHW`. |
361 | |
362 | )code" TVM_ADD_FILELINE) |
363 | .set_attrs_type<GlobalPool2DAttrs>() |
364 | .set_num_inputs(1) |
365 | .add_argument("data" , "Tensor" , "The input tensor." ) |
366 | .set_support_level(2) |
367 | .add_type_rel("GlobalMaxPool2D" , GlobalPool2DRel) |
368 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , PoolInferCorrectLayout<GlobalPool2DAttrs>) |
369 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
370 | .set_attr<FTVMCompute>("FTVMCompute" , GlobalPool2DCompute<topi::nn::kMaxPool>); |
371 | |
372 | // relay.nn.adaptive_pool_1d |
373 | TVM_REGISTER_NODE_TYPE(AdaptivePool1DAttrs); |
374 | |
375 | bool AdaptivePool1DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, |
376 | const TypeReporter& reporter) { |
377 | ICHECK_EQ(types.size(), 2); |
378 | const auto* data = types[0].as<TensorTypeNode>(); |
379 | if (data == nullptr) { |
380 | return false; |
381 | } |
382 | const auto dshape = data->shape; |
383 | ICHECK_GE(dshape.size(), 1U) << "Pool2D only support input >= 1-D: input must have width" ; |
384 | const auto* param = attrs.as<AdaptivePool1DAttrs>(); |
385 | ICHECK(param != nullptr); |
386 | |
387 | Layout layout(param->layout); |
388 | ICHECK(layout.Contains(LayoutAxis::Get('W')) && !layout.Contains(LayoutAxis::Get('w'))) |
389 | << "Invalid layout " << layout << ". Pool1D layout must have W, which cannot be split" ; |
390 | |
391 | const auto widx = layout.IndexOf(LayoutAxis::Get('W')); |
392 | Array<IndexExpr> oshape(dshape); |
393 | auto output_size = param->output_size; |
394 | ICHECK_LE(output_size.size(), 1U) << "output_size must have 1 element." ; |
395 | IndexExpr output_width; |
396 | if (output_size.empty()) { |
397 | output_width = dshape[widx]; |
398 | } else { |
399 | output_width = output_size[0]; |
400 | } |
401 | |
402 | oshape.Set(widx, output_width); |
403 | |
404 | // assign output type |
405 | reporter->Assign(types[1], TensorType(oshape, data->dtype)); |
406 | return true; |
407 | } |
408 | |
409 | template <topi::nn::PoolType mode> |
410 | Array<te::Tensor> AdaptivePool1DCompute(const Attrs& attrs, const Array<te::Tensor>& inputs, |
411 | const Type& out_type) { |
412 | static const Layout kNCW("NCW" ); |
413 | const auto* param = attrs.as<AdaptivePool1DAttrs>(); |
414 | ICHECK(param != nullptr); |
415 | Layout layout(param->layout); |
416 | ICHECK(tir::BijectiveLayout(layout, kNCW).defined()) |
417 | << "Adaptive pool1d currently only supports layouts that are convertible from NCW" ; |
418 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('w')), -1) |
419 | << "Adaptive pool2d does not support input split on width" ; |
420 | |
421 | ICHECK(inputs[0].ndim() == 3U || inputs[0].ndim() == 4U) |
422 | << "Pool1D only support 3-D input (e.g., NCW)" |
423 | << " or 4-D input (last dimension is a split of channel)" ; |
424 | |
425 | auto output_size = param->output_size; |
426 | const auto widx = layout.IndexOf(LayoutAxis::Get('W')); |
427 | IndexExpr output_width; |
428 | if (output_size.empty()) { |
429 | output_width = inputs[0]->shape[widx]; |
430 | } else { |
431 | output_width = output_size[0]; |
432 | } |
433 | return Array<te::Tensor>{ |
434 | topi::nn::adaptive_pool1d(inputs[0], Array<IndexExpr>{output_width}, mode, layout.name())}; |
435 | } |
436 | |
437 | // relay.nn.adaptive_avg_pool1d |
438 | Expr MakeAdaptiveAvgPool1D(Expr data, Array<IndexExpr> output_size, String layout, |
439 | String out_layout) { |
440 | auto attrs = make_object<AdaptivePool1DAttrs>(); |
441 | attrs->output_size = std::move(output_size); |
442 | attrs->layout = std::move(layout); |
443 | attrs->out_layout = std::move(out_layout); |
444 | static const Op& op = Op::Get("nn.adaptive_avg_pool1d" ); |
445 | return Call(op, {data}, Attrs(attrs), {}); |
446 | } |
447 | |
448 | TVM_REGISTER_GLOBAL("relay.op.nn._make.adaptive_avg_pool1d" ).set_body_typed(MakeAdaptiveAvgPool1D); |
449 | |
450 | RELAY_REGISTER_OP("nn.adaptive_avg_pool1d" ) |
451 | .describe(R"code(Adaptive average pooling operation for 1D data. |
452 | |
453 | - **data**: This depends on the `layout` parameter. Input is 3D array of shape |
454 | (batch_size, channels, width) if `layout` is `NCW`. |
455 | - **output_size**: If this argument is not provided, input width will be used |
456 | as output width. |
457 | If an integer is provided for output_size, the output size is |
458 | (N x C x output_size) for any input (NCW). |
459 | - **out**: This depends on the `layout` parameter. Output is 3D array of shape |
460 | (batch_size, channels, output_width) if `layout` is `NCW`. |
461 | |
462 | )code" TVM_ADD_FILELINE) |
463 | .set_attrs_type<AdaptivePool1DAttrs>() |
464 | .set_num_inputs(1) |
465 | .add_argument("data" , "Tensor" , "The input tensor." ) |
466 | .set_support_level(10) |
467 | .add_type_rel("AdaptiveAvgPool1D" , AdaptivePool1DRel) |
468 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , |
469 | PoolInferCorrectLayout<AdaptivePool1DAttrs>) |
470 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
471 | .set_attr<FTVMCompute>("FTVMCompute" , AdaptivePool1DCompute<topi::nn::kAvgPool>); |
472 | |
473 | // relay.nn.adaptive_max_pool1d |
474 | Expr MakeAdaptiveMaxPool1D(Expr data, Array<IndexExpr> output_size, String layout, |
475 | String out_layout) { |
476 | auto attrs = make_object<AdaptivePool1DAttrs>(); |
477 | attrs->output_size = std::move(output_size); |
478 | attrs->layout = std::move(layout); |
479 | attrs->out_layout = std::move(out_layout); |
480 | static const Op& op = Op::Get("nn.adaptive_max_pool1d" ); |
481 | return Call(op, {data}, Attrs(attrs), {}); |
482 | } |
483 | |
484 | TVM_REGISTER_GLOBAL("relay.op.nn._make.adaptive_max_pool1d" ).set_body_typed(MakeAdaptiveMaxPool1D); |
485 | |
486 | RELAY_REGISTER_OP("nn.adaptive_max_pool1d" ) |
487 | .describe(R"code(Adaptive max pooling operation for 1D data. |
488 | |
489 | - **data**: This depends on the `layout` parameter. Input is 3D array of shape |
490 | (batch_size, channels, width) if `layout` is `NCW`. |
491 | - **output_size**: If this argument is not provided, input width will be used |
492 | as output width. |
493 | If an integer is provided for output_size, the output size is |
494 | (N x C x output_size) for any input (NCW). |
495 | - **out**: This depends on the `layout` parameter. Output is 3D array of shape |
496 | (batch_size, channels, output_width) if `layout` is `NCW`. |
497 | |
498 | )code" TVM_ADD_FILELINE) |
499 | .set_attrs_type<AdaptivePool1DAttrs>() |
500 | .set_num_inputs(1) |
501 | .add_argument("data" , "Tensor" , "The input tensor." ) |
502 | .set_support_level(10) |
503 | .add_type_rel("AdaptiveMaxPool1D" , AdaptivePool1DRel) |
504 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , |
505 | PoolInferCorrectLayout<AdaptivePool1DAttrs>) |
506 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
507 | .set_attr<FTVMCompute>("FTVMCompute" , AdaptivePool1DCompute<topi::nn::kMaxPool>); |
508 | |
509 | // relay.nn.adaptive_pool_2d |
510 | TVM_REGISTER_NODE_TYPE(AdaptivePool2DAttrs); |
511 | |
512 | bool AdaptivePool2DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, |
513 | const TypeReporter& reporter) { |
514 | ICHECK_EQ(types.size(), 2); |
515 | const auto* data = types[0].as<TensorTypeNode>(); |
516 | if (data == nullptr) { |
517 | return false; |
518 | } |
519 | const auto dshape = data->shape; |
520 | ICHECK_GE(dshape.size(), 2U) |
521 | << "Pool2D only support input >= 2-D: input must have height and width" ; |
522 | const auto* param = attrs.as<AdaptivePool2DAttrs>(); |
523 | ICHECK(param != nullptr); |
524 | |
525 | Layout layout(param->layout); |
526 | ICHECK(layout.Contains(LayoutAxis::Get('H')) && layout.Contains(LayoutAxis::Get('W')) && |
527 | !layout.Contains(LayoutAxis::Get('h')) && !layout.Contains(LayoutAxis::Get('w'))) |
528 | << "Invalid layout " << layout << ". Pool2D layout must have H and W, which cannot be split" ; |
529 | |
530 | const auto hidx = layout.IndexOf(LayoutAxis::Get('H')); |
531 | const auto widx = layout.IndexOf(LayoutAxis::Get('W')); |
532 | Array<IndexExpr> oshape(dshape); |
533 | auto output_size = param->output_size; |
534 | ICHECK_LE(output_size.size(), 2U) << "output_size can have up to 2 elements." ; |
535 | IndexExpr output_height, output_width; |
536 | if (output_size.empty()) { |
537 | output_height = dshape[hidx]; |
538 | output_width = dshape[widx]; |
539 | } else if (output_size.size() == 1) { |
540 | output_height = output_size[0]; |
541 | output_width = output_size[0]; |
542 | } else { |
543 | output_height = output_size[0]; |
544 | output_width = output_size[1]; |
545 | } |
546 | |
547 | oshape.Set(hidx, output_height); |
548 | oshape.Set(widx, output_width); |
549 | |
550 | // assign output type |
551 | reporter->Assign(types[1], TensorType(oshape, data->dtype)); |
552 | return true; |
553 | } |
554 | |
555 | template <topi::nn::PoolType mode> |
556 | Array<te::Tensor> AdaptivePool2DCompute(const Attrs& attrs, const Array<te::Tensor>& inputs, |
557 | const Type& out_type) { |
558 | static const Layout kNCHW("NCHW" ); |
559 | const auto* param = attrs.as<AdaptivePool2DAttrs>(); |
560 | ICHECK(param != nullptr); |
561 | Layout layout(param->layout); |
562 | ICHECK(tir::BijectiveLayout(layout, kNCHW).defined()) |
563 | << "Adaptive pool2d currently only supports layouts that are convertible from NCHW" ; |
564 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('h')), -1) |
565 | << "Adaptive pool2d does not support input split on height" ; |
566 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('w')), -1) |
567 | << "Adaptive pool2d does not support input split on width" ; |
568 | |
569 | ICHECK(inputs[0].ndim() == 4U || inputs[0].ndim() == 5U) |
570 | << "Pool2D only support 4-D input (e.g., NCHW)" |
571 | << " or 5-D input (last dimension is a split of channel)" ; |
572 | |
573 | auto output_size = param->output_size; |
574 | const auto hidx = layout.IndexOf(LayoutAxis::Get('H')); |
575 | const auto widx = layout.IndexOf(LayoutAxis::Get('W')); |
576 | IndexExpr output_height, output_width; |
577 | if (output_size.empty()) { |
578 | output_height = inputs[0]->shape[hidx]; |
579 | output_width = inputs[0]->shape[widx]; |
580 | } else if (output_size.size() == 1) { |
581 | output_height = output_size[0]; |
582 | output_width = output_size[0]; |
583 | } else { |
584 | output_height = output_size[0]; |
585 | output_width = output_size[1]; |
586 | } |
587 | return Array<te::Tensor>{topi::nn::adaptive_pool( |
588 | inputs[0], Array<IndexExpr>{output_height, output_width}, mode, layout.name())}; |
589 | } |
590 | |
591 | // relay.nn.adaptive_avg_pool2d |
592 | Expr MakeAdaptiveAvgPool2D(Expr data, Array<IndexExpr> output_size, String layout, |
593 | String out_layout) { |
594 | auto attrs = make_object<AdaptivePool2DAttrs>(); |
595 | attrs->output_size = std::move(output_size); |
596 | attrs->layout = std::move(layout); |
597 | attrs->out_layout = std::move(out_layout); |
598 | static const Op& op = Op::Get("nn.adaptive_avg_pool2d" ); |
599 | return Call(op, {data}, Attrs(attrs), {}); |
600 | } |
601 | |
602 | TVM_REGISTER_GLOBAL("relay.op.nn._make.adaptive_avg_pool2d" ).set_body_typed(MakeAdaptiveAvgPool2D); |
603 | |
604 | RELAY_REGISTER_OP("nn.adaptive_avg_pool2d" ) |
605 | .describe(R"code(Adaptive average pooling operation for 2D data. |
606 | |
607 | - **data**: This depends on the `layout` parameter. Input is 4D array of shape |
608 | (batch_size, channels, height, width) if `layout` is `NCHW`. |
609 | - **output_size**: If this argument is not provided, input height and width will be used |
610 | as output height and width. |
611 | If a single integer is provided for output_size, the output size is |
612 | (N x C x output_size x output_size) for any input (NCHW). |
613 | If a tuple of integers (height, width) are provided for output_size, |
614 | the output size is (N x C x height x width) for any input (NCHW). |
615 | - **out**: This depends on the `layout` parameter. Output is 4D array of shape |
616 | (batch_size, channels, output_height, output_width) if `layout` is `NCHW`. |
617 | |
618 | )code" TVM_ADD_FILELINE) |
619 | .set_attrs_type<AdaptivePool2DAttrs>() |
620 | .set_num_inputs(1) |
621 | .add_argument("data" , "Tensor" , "The input tensor." ) |
622 | .set_support_level(10) |
623 | .add_type_rel("AdaptiveAvgPool2D" , AdaptivePool2DRel) |
624 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , |
625 | PoolInferCorrectLayout<AdaptivePool2DAttrs>) |
626 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
627 | .set_attr<FTVMCompute>("FTVMCompute" , AdaptivePool2DCompute<topi::nn::kAvgPool>); |
628 | |
629 | // relay.nn.adaptive_max_pool2d |
630 | Expr MakeAdaptiveMaxPool2D(Expr data, Array<IndexExpr> output_size, String layout, |
631 | String out_layout) { |
632 | auto attrs = make_object<AdaptivePool2DAttrs>(); |
633 | attrs->output_size = std::move(output_size); |
634 | attrs->layout = std::move(layout); |
635 | attrs->out_layout = std::move(out_layout); |
636 | static const Op& op = Op::Get("nn.adaptive_max_pool2d" ); |
637 | return Call(op, {data}, Attrs(attrs), {}); |
638 | } |
639 | |
640 | TVM_REGISTER_GLOBAL("relay.op.nn._make.adaptive_max_pool2d" ).set_body_typed(MakeAdaptiveMaxPool2D); |
641 | |
642 | RELAY_REGISTER_OP("nn.adaptive_max_pool2d" ) |
643 | .describe(R"code(Adaptive max pooling operation for 2D data. |
644 | |
645 | - **data**: This depends on the `layout` parameter. Input is 4D array of shape |
646 | (batch_size, channels, height, width) if `layout` is `NCHW`. |
647 | - **output_size**: If this argument is not provided, input height and width will be used |
648 | as output height and width. |
649 | If a single integer is provided for output_size, the output size is |
650 | (N x C x output_size x output_size) for any input (NCHW). |
651 | If a tuple of integers (height, width) are provided for output_size, |
652 | the output size is (N x C x height x width) for any input (NCHW). |
653 | - **out**: This depends on the `layout` parameter. Output is 4D array of shape |
654 | (batch_size, channels, output_height, output_width) if `layout` is `NCHW`. |
655 | |
656 | )code" TVM_ADD_FILELINE) |
657 | .set_attrs_type<AdaptivePool2DAttrs>() |
658 | .set_num_inputs(1) |
659 | .add_argument("data" , "Tensor" , "The input tensor." ) |
660 | .set_support_level(10) |
661 | .add_type_rel("AdaptiveMaxPool2D" , AdaptivePool2DRel) |
662 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , |
663 | PoolInferCorrectLayout<AdaptivePool2DAttrs>) |
664 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
665 | .set_attr<FTVMCompute>("FTVMCompute" , AdaptivePool2DCompute<topi::nn::kMaxPool>); |
666 | |
667 | // relay.nn.adaptive_pool3d |
668 | TVM_REGISTER_NODE_TYPE(AdaptivePool3DAttrs); |
669 | |
670 | bool AdaptivePool3DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, |
671 | const TypeReporter& reporter) { |
672 | ICHECK_EQ(types.size(), 2); |
673 | const auto* data = types[0].as<TensorTypeNode>(); |
674 | if (data == nullptr) { |
675 | return false; |
676 | } |
677 | const auto dshape = data->shape; |
678 | ICHECK_GE(dshape.size(), 3U) |
679 | << "Pool3D only support input >= 3-D: input must have depth, height and width" ; |
680 | const auto* param = attrs.as<AdaptivePool3DAttrs>(); |
681 | ICHECK(param != nullptr); |
682 | |
683 | Layout layout(param->layout); |
684 | ICHECK(layout.Contains(LayoutAxis::Get('D')) && layout.Contains(LayoutAxis::Get('H')) && |
685 | layout.Contains(LayoutAxis::Get('W')) && !layout.Contains(LayoutAxis::Get('d')) && |
686 | !layout.Contains(LayoutAxis::Get('h')) && !layout.Contains(LayoutAxis::Get('w'))) |
687 | << "Invalid layout " << layout |
688 | << ". Pool3D layout must have D, H and W, which cannot be split" ; |
689 | |
690 | const auto didx = layout.IndexOf(LayoutAxis::Get('D')); |
691 | const auto hidx = layout.IndexOf(LayoutAxis::Get('H')); |
692 | const auto widx = layout.IndexOf(LayoutAxis::Get('W')); |
693 | Array<IndexExpr> oshape(dshape); |
694 | auto output_size = param->output_size; |
695 | ICHECK_LE(output_size.size(), 3U) << "output_size can have up to 3 elements." ; |
696 | IndexExpr output_depth, output_height, output_width; |
697 | if (output_size.empty()) { |
698 | output_depth = dshape[didx]; |
699 | output_height = dshape[hidx]; |
700 | output_width = dshape[widx]; |
701 | } else if (output_size.size() == 1) { |
702 | output_depth = output_size[0]; |
703 | output_height = output_size[0]; |
704 | output_width = output_size[0]; |
705 | } else { |
706 | output_depth = output_size[0]; |
707 | output_height = output_size[1]; |
708 | output_width = output_size[2]; |
709 | } |
710 | |
711 | oshape.Set(didx, output_depth); |
712 | oshape.Set(hidx, output_height); |
713 | oshape.Set(widx, output_width); |
714 | |
715 | // assign output type |
716 | reporter->Assign(types[1], TensorType(oshape, data->dtype)); |
717 | return true; |
718 | } |
719 | |
720 | template <topi::nn::PoolType mode> |
721 | Array<te::Tensor> AdaptivePool3DCompute(const Attrs& attrs, const Array<te::Tensor>& inputs, |
722 | const Type& out_type) { |
723 | static const Layout kNCDHW("NCDHW" ); |
724 | const auto* param = attrs.as<AdaptivePool3DAttrs>(); |
725 | ICHECK(param != nullptr); |
726 | Layout layout(param->layout); |
727 | Layout out_layout(param->out_layout); |
728 | ICHECK(tir::BijectiveLayout(layout, kNCDHW).defined()) |
729 | << "Adaptive pool3d currently only supports layouts that are convertible from NCDHW" ; |
730 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('d')), -1) |
731 | << "Adaptive pool3d does not support input split on depth" ; |
732 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('h')), -1) |
733 | << "Adaptive pool3d does not support input split on height" ; |
734 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('w')), -1) |
735 | << "Adaptive pool3d does not support input split on width" ; |
736 | |
737 | ICHECK(inputs[0].ndim() == 5U || inputs[0].ndim() == 6U) |
738 | << "Pool3D only support 5-D input (e.g., NCDHW)" |
739 | << " or 6-D input (last dimension is a split of channel)" ; |
740 | |
741 | auto output_size = param->output_size; |
742 | const auto didx = layout.IndexOf(LayoutAxis::Get('D')); |
743 | const auto hidx = layout.IndexOf(LayoutAxis::Get('H')); |
744 | const auto widx = layout.IndexOf(LayoutAxis::Get('W')); |
745 | IndexExpr output_depth, output_height, output_width; |
746 | if (output_size.empty()) { |
747 | output_depth = inputs[0]->shape[didx]; |
748 | output_height = inputs[0]->shape[hidx]; |
749 | output_width = inputs[0]->shape[widx]; |
750 | } else if (output_size.size() == 1) { |
751 | output_depth = output_size[0]; |
752 | output_height = output_size[0]; |
753 | output_width = output_size[0]; |
754 | } else { |
755 | output_depth = output_size[0]; |
756 | output_height = output_size[1]; |
757 | output_width = output_size[2]; |
758 | } |
759 | |
760 | auto osize = Array<IndexExpr>{output_depth, output_height, output_width}; |
761 | return Array<te::Tensor>{topi::nn::adaptive_pool3d(inputs[0], osize, mode, layout.name())}; |
762 | } |
763 | |
764 | // relay.nn.adaptive_max_pool3d |
765 | Expr MakeAdaptiveMaxPool3D(Expr data, Array<IndexExpr> output_size, String layout, |
766 | String out_layout) { |
767 | auto attrs = make_object<AdaptivePool3DAttrs>(); |
768 | attrs->output_size = std::move(output_size); |
769 | attrs->layout = std::move(layout); |
770 | attrs->out_layout = std::move(out_layout); |
771 | static const Op& op = Op::Get("nn.adaptive_max_pool3d" ); |
772 | return Call(op, {data}, Attrs(attrs), {}); |
773 | } |
774 | |
775 | TVM_REGISTER_GLOBAL("relay.op.nn._make.adaptive_max_pool3d" ).set_body_typed(MakeAdaptiveMaxPool3D); |
776 | |
777 | RELAY_REGISTER_OP("nn.adaptive_max_pool3d" ) |
778 | .describe(R"code(Adaptive max pooling operation for 3D data. |
779 | |
780 | - **data**: This depends on the `layout` parameter. Input is 5D array of shape |
781 | (batch_size, channels, depth, height, width) if `layout` is `NCDHW`. |
782 | - **output_size**: If this argument is not provided, input depth, height and width will be used |
783 | as output depth, height and width. |
784 | If a single integer is provided for output_size, the output size is |
785 | (N x C x output_size x output_size x output_size) for any input (NCDHW). |
786 | If a tuple of integers (depth, height, width) are provided for output_size, |
787 | the output size is (N x C x depth x height x width) for any input (NCDHW). |
788 | - **out**: This depends on the `layout` parameter. Output is 5D array of shape |
789 | (batch_size, channels, output_depth, output_height, output_width) if `layout` is `NCDHW`. |
790 | |
791 | )code" TVM_ADD_FILELINE) |
792 | .set_attrs_type<AdaptivePool3DAttrs>() |
793 | .set_num_inputs(1) |
794 | .add_argument("data" , "Tensor" , "The input tensor." ) |
795 | .set_support_level(10) |
796 | .add_type_rel("AdaptiveMaxPool3D" , AdaptivePool3DRel) |
797 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , |
798 | PoolInferCorrectLayout<AdaptivePool3DAttrs>) |
799 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
800 | .set_attr<FTVMCompute>("FTVMCompute" , AdaptivePool3DCompute<topi::nn::kMaxPool>); |
801 | |
802 | // relay.nn.adaptive_max_pool3d |
803 | Expr MakeAdaptiveAvgPool3D(Expr data, Array<IndexExpr> output_size, String layout, |
804 | String out_layout) { |
805 | auto attrs = make_object<AdaptivePool3DAttrs>(); |
806 | attrs->output_size = std::move(output_size); |
807 | attrs->layout = std::move(layout); |
808 | attrs->out_layout = std::move(out_layout); |
809 | static const Op& op = Op::Get("nn.adaptive_avg_pool3d" ); |
810 | return Call(op, {data}, Attrs(attrs), {}); |
811 | } |
812 | |
813 | TVM_REGISTER_GLOBAL("relay.op.nn._make.adaptive_avg_pool3d" ).set_body_typed(MakeAdaptiveAvgPool3D); |
814 | |
815 | RELAY_REGISTER_OP("nn.adaptive_avg_pool3d" ) |
816 | .describe(R"code(Adaptive avg pooling operation for 3D data. |
817 | - **data**: This depends on the `layout` parameter. Input is 5D array of shape |
818 | (batch_size, channels, depth, height, width) if `layout` is `NCDHW`. |
819 | - **output_size**: If this argument is not provided, input depth, height and width will be used |
820 | as output depth, height and width. |
821 | If a single integer is provided for output_size, the output size is |
822 | (N x C x output_size x output_size x output_size) for any input (NCDHW). |
823 | If a tuple of integers (depth, height, width) are provided for output_size, |
824 | the output size is (N x C x depth x height x width) for any input (NCDHW). |
825 | - **out**: This depends on the `layout` parameter. Output is 5D array of shape |
826 | (batch_size, channels, output_depth, output_height, output_width) if `layout` is `NCDHW`. |
827 | )code" TVM_ADD_FILELINE) |
828 | .set_attrs_type<AdaptivePool3DAttrs>() |
829 | .set_num_inputs(1) |
830 | .add_argument("data" , "Tensor" , "The input tensor." ) |
831 | .set_support_level(10) |
832 | .add_type_rel("AdaptiveAvgPool3D" , AdaptivePool3DRel) |
833 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , |
834 | PoolInferCorrectLayout<AdaptivePool3DAttrs>) |
835 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
836 | .set_attr<FTVMCompute>("FTVMCompute" , AdaptivePool3DCompute<topi::nn::kAvgPool>); |
837 | |
838 | bool Pool2DGradRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, |
839 | const TypeReporter& reporter) { |
840 | ICHECK_EQ(types.size(), 3); |
841 | const auto* data = types[1].as<TensorTypeNode>(); |
842 | |
843 | if (data == nullptr) return false; |
844 | |
845 | // assign output type |
846 | reporter->Assign(types[2], types[1]); |
847 | return true; |
848 | } |
849 | |
850 | template <typename AttrType, topi::nn::PoolType mode> |
851 | Array<te::Tensor> Pool2DGradCompute(const Attrs& attrs, const Array<te::Tensor>& inputs, |
852 | const Type& out_type) { |
853 | static const Layout kNCHW("NCHW" ); |
854 | const auto* param = attrs.as<AttrType>(); |
855 | ICHECK(param != nullptr); |
856 | ICHECK_EQ(inputs.size(), 2); |
857 | auto pool_size = param->pool_size; |
858 | auto strides = param->strides; |
859 | auto padding = param->padding; |
860 | auto ceil_mode = param->ceil_mode; |
861 | Layout layout(param->layout); |
862 | |
863 | ICHECK(tir::BijectiveLayout(layout, kNCHW).defined()) |
864 | << "pool2d_grad currently only supports layouts that are convertible from NCHW" ; |
865 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('h')), -1) |
866 | << "pool2d_grad does not support input split on height" ; |
867 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('w')), -1) |
868 | << "pool2d_grad does not support input split on width" ; |
869 | |
870 | ICHECK(inputs[0].ndim() == 4U || inputs[0].ndim() == 5U) |
871 | << "Pool2DGrad only support 4-D output gradient (e.g., NCHW)" |
872 | << " or 5-D output gradient (last dimension is a split of channel)" ; |
873 | |
874 | ICHECK(inputs[1].ndim() == 4U || inputs[1].ndim() == 5U) |
875 | << "Pool2DGrad only support 4-D input (e.g., NCHW)" |
876 | << " or 5-D input (last dimension is a split of channel)" ; |
877 | |
878 | if (param->padding.size() == 1) { |
879 | padding.push_back(padding[0]); |
880 | padding.push_back(padding[0]); |
881 | padding.push_back(padding[0]); |
882 | } else if (param->padding.size() == 2) { |
883 | padding.push_back(padding[0]); |
884 | padding.push_back(padding[1]); |
885 | } |
886 | if (mode == topi::nn::kAvgPool) { |
887 | bool count_include_pad = reinterpret_cast<const AvgPool2DAttrs*>(param)->count_include_pad; |
888 | return Array<te::Tensor>{topi::nn::pool_grad(inputs[0], inputs[1], pool_size, strides, padding, |
889 | mode, ceil_mode, layout.name(), |
890 | count_include_pad)}; |
891 | } else { |
892 | return Array<te::Tensor>{topi::nn::pool_grad(inputs[0], inputs[1], pool_size, strides, padding, |
893 | mode, ceil_mode, layout.name())}; |
894 | } |
895 | } |
896 | |
897 | // MaxPool2DGrad |
898 | Expr MakeMaxPool2DGrad(Expr out_grad, Expr data, Array<IndexExpr> pool_size, |
899 | Array<IndexExpr> strides, Array<IndexExpr> padding, String layout, |
900 | String out_layout, bool ceil_mode) { |
901 | auto attrs = make_object<MaxPool2DAttrs>(); |
902 | attrs->pool_size = std::move(pool_size); |
903 | attrs->strides = std::move(strides); |
904 | attrs->padding = std::move(padding); |
905 | attrs->layout = std::move(layout); |
906 | attrs->out_layout = std::move(out_layout); |
907 | attrs->ceil_mode = ceil_mode; |
908 | static const Op& op = Op::Get("nn.max_pool2d_grad" ); |
909 | return Call(op, {out_grad, data}, Attrs(attrs), {}); |
910 | } |
911 | |
912 | TVM_REGISTER_GLOBAL("relay.op.nn._make.max_pool2d_grad" ).set_body_typed(MakeMaxPool2DGrad); |
913 | |
914 | RELAY_REGISTER_OP("nn.max_pool2d_grad" ) |
915 | .describe(R"code(Gradient of max pooling operation for two dimensional data. |
916 | |
917 | - **out_grad**: This depends on the `layout` parameter. Output gradient is 4D array of |
918 | shape (batch_size, channels, out_height, out_width) if `layout` is `NCHW`. |
919 | out_height and out_width are the output size of the pooling operation, |
920 | which are calculated as:: |
921 | out_height = floor((height+padding[0]+padding[2]-pool_size[0])/strides[0])+1 |
922 | out_width = floor((width+padding[1]+padding[3]-pool_size[1])/strides[1])+1 |
923 | |
924 | where padding will be an expanded array based on number of values passed as:: |
925 | one int : all sides same padding used. |
926 | two int : bottom, right use same as top and left. |
927 | four int: padding width in the order of (top, left, bottom, right). |
928 | |
929 | When `ceil_mode` is `True`, ceil will be used instead of floor in this |
930 | equation. |
931 | - **data**: This depends on the `layout` parameter. Input is 4D array of shape |
932 | (batch_size, channels, height, width) if `layout` is `NCHW`. |
933 | - **grad**: This depends on the `layout` parameter. Grad is 4D array of shape |
934 | (batch_size, channels, height, width) if `layout` is `NCHW`. |
935 | |
936 | )code" TVM_ADD_FILELINE) |
937 | .set_attrs_type<MaxPool2DAttrs>() |
938 | .set_num_inputs(2) |
939 | .add_argument("data" , "Tensor" , "The input tensor." ) |
940 | .add_argument("grad" , "Tensor" , "The grad tensor." ) |
941 | .set_support_level(2) |
942 | .add_type_rel("MaxPool2DGrad" , Pool2DGradRel) |
943 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
944 | .set_attr<FTVMCompute>("FTVMCompute" , Pool2DGradCompute<MaxPool2DAttrs, topi::nn::kMaxPool>); |
945 | |
946 | // AvgPool2DGrad |
947 | Expr MakeAvgPool2DGrad(Expr out_grad, Expr data, Array<IndexExpr> pool_size, |
948 | Array<IndexExpr> strides, Array<IndexExpr> padding, String layout, |
949 | String out_layout, bool ceil_mode, bool count_include_pad) { |
950 | auto attrs = make_object<AvgPool2DAttrs>(); |
951 | attrs->pool_size = std::move(pool_size); |
952 | attrs->strides = std::move(strides); |
953 | attrs->padding = std::move(padding); |
954 | attrs->layout = std::move(layout); |
955 | attrs->out_layout = std::move(out_layout); |
956 | attrs->ceil_mode = ceil_mode; |
957 | attrs->count_include_pad = count_include_pad; |
958 | static const Op& op = Op::Get("nn.avg_pool2d_grad" ); |
959 | return Call(op, {out_grad, data}, Attrs(attrs), {}); |
960 | } |
961 | |
962 | TVM_REGISTER_GLOBAL("relay.op.nn._make.avg_pool2d_grad" ).set_body_typed(MakeAvgPool2DGrad); |
963 | |
964 | RELAY_REGISTER_OP("nn.avg_pool2d_grad" ) |
965 | .describe(R"code(Gradient of average pooling operation for two dimensional data. |
966 | |
967 | - **out_grad**: This depends on the `layout` parameter. Output gradient is 4D array of |
968 | shape (batch_size, channels, out_height, out_width) if `layout` is `NCHW`. |
969 | out_height and out_width are the output size of the pooling operation, |
970 | which are calculated as:: |
971 | out_height = floor((height+padding[0]+padding[2]-pool_size[0])/strides[0])+1 |
972 | out_width = floor((width+padding[1]+padding[3]-pool_size[1])/strides[1])+1 |
973 | |
974 | where padding will be an expanded array based on number of values passed as:: |
975 | one int : all sides same padding used. |
976 | two int : bottom, right use same as top and left. |
977 | four int: padding width in the order of (top, left, bottom, right). |
978 | |
979 | When `ceil_mode` is `True`, ceil will be used instead of floor in this |
980 | equation. |
981 | - **data**: This depends on the `layout` parameter. Input is 4D array of shape |
982 | (batch_size, channels, height, width) if `layout` is `NCHW`. |
983 | - **grad**: This depends on the `layout` parameter. Grad is 4D array of shape |
984 | (batch_size, channels, height, width) if `layout` is `NCHW`. |
985 | |
986 | )code" TVM_ADD_FILELINE) |
987 | .set_attrs_type<MaxPool2DAttrs>() |
988 | .set_num_inputs(2) |
989 | .add_argument("data" , "Tensor" , "The input tensor." ) |
990 | .add_argument("grad" , "Tensor" , "The grad tensor." ) |
991 | .set_support_level(2) |
992 | .add_type_rel("MaxPool2DGrad" , Pool2DGradRel) |
993 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
994 | .set_attr<FTVMCompute>("FTVMCompute" , Pool2DGradCompute<AvgPool2DAttrs, topi::nn::kAvgPool>); |
995 | |
996 | // relay.nn.max_pool1d & relay.nn.avg_pool1d |
997 | TVM_REGISTER_NODE_TYPE(MaxPool1DAttrs); |
998 | TVM_REGISTER_NODE_TYPE(AvgPool1DAttrs); |
999 | |
1000 | template <typename AttrType> |
1001 | bool Pool1DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, |
1002 | const TypeReporter& reporter) { |
1003 | ICHECK_EQ(types.size(), 2); |
1004 | const auto* data = types[0].as<TensorTypeNode>(); |
1005 | |
1006 | if (data == nullptr) return false; |
1007 | |
1008 | const auto dshape = data->shape; |
1009 | ICHECK_GE(dshape.size(), 1U) << "Pool1D only support input >= 1-D: input must have width" ; |
1010 | const auto param = attrs.as<AttrType>(); |
1011 | ICHECK(param != nullptr); |
1012 | |
1013 | Layout layout(param->layout); |
1014 | Layout out_layout(param->out_layout); |
1015 | ICHECK(layout.Contains(LayoutAxis::Get('W')) && !layout.Contains(LayoutAxis::Get('w'))) |
1016 | << "Invalid layout " << layout << ". Pool1D layout must have W, which cannot be split" ; |
1017 | |
1018 | const auto widx = layout.IndexOf(LayoutAxis::Get('W')); |
1019 | |
1020 | IndexExpr pad_w; |
1021 | if (param->padding.size() == 1) { |
1022 | pad_w = param->padding[0] * 2; |
1023 | } else if (param->padding.size() == 2) { |
1024 | // (left, right) |
1025 | pad_w = param->padding[0] + param->padding[1]; |
1026 | } else { |
1027 | return false; |
1028 | } |
1029 | |
1030 | std::vector<IndexExpr> oshape(dshape.begin(), dshape.end()); |
1031 | |
1032 | if (dshape[widx].as<tir::AnyNode>()) { |
1033 | oshape[widx] = dshape[widx]; |
1034 | } else { |
1035 | oshape[widx] = |
1036 | calculate_pool_dimension(dshape[widx], pad_w, param->pool_size[0], param->dilation[0], |
1037 | param->strides[0], param->ceil_mode); |
1038 | } |
1039 | |
1040 | // assign output type |
1041 | reporter->Assign(types[1], TensorType(oshape, data->dtype)); |
1042 | return true; |
1043 | } |
1044 | |
1045 | template <typename AttrType, topi::nn::PoolType mode> |
1046 | Array<te::Tensor> Pool1DCompute(const Attrs& attrs, const Array<te::Tensor>& inputs, |
1047 | const Type& out_type) { |
1048 | static const Layout kNCW("NCW" ); |
1049 | const auto* param = attrs.as<AttrType>(); |
1050 | ICHECK(param != nullptr); |
1051 | auto pool_size = param->pool_size; |
1052 | auto strides = param->strides; |
1053 | auto dilation = param->dilation; |
1054 | auto padding = param->padding; |
1055 | auto ceil_mode = param->ceil_mode; |
1056 | Layout layout(param->layout); |
1057 | Layout out_layout(param->out_layout); |
1058 | |
1059 | ICHECK(tir::BijectiveLayout(layout, kNCW).defined()) |
1060 | << "max_pool1d currently only supports layouts that are convertible from NCW" ; |
1061 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('w')), -1) |
1062 | << "max_pool1d does not support input split on width" ; |
1063 | |
1064 | ICHECK(inputs[0].ndim() == 3U || inputs[0].ndim() == 4U || inputs[0].ndim() == 5U) |
1065 | << "Pool1D only support 3-D input (e.g., NCW)" |
1066 | << " or 4-D input (e.g. NCWc on for vector instructions)" |
1067 | << " or 5-D input (e.g. NCWnc for tensor accelerators)" ; |
1068 | |
1069 | if (param->padding.size() == 1) { |
1070 | padding.push_back(padding[0]); |
1071 | } |
1072 | |
1073 | if (mode == topi::nn::kAvgPool) { |
1074 | bool count_include_pad = reinterpret_cast<const AvgPool1DAttrs*>(param)->count_include_pad; |
1075 | return Array<te::Tensor>{topi::nn::pool1d(inputs[0], pool_size, strides, dilation, padding, |
1076 | mode, ceil_mode, layout.name(), count_include_pad)}; |
1077 | } else { |
1078 | return Array<te::Tensor>{topi::nn::pool1d(inputs[0], pool_size, strides, dilation, padding, |
1079 | mode, ceil_mode, layout.name())}; |
1080 | } |
1081 | } |
1082 | |
1083 | TVM_REGISTER_GLOBAL("relay.op.nn._make.max_pool1d" ) |
1084 | .set_body_typed([](Expr data, Array<IndexExpr> pool_size, Array<IndexExpr> strides, |
1085 | Array<IndexExpr> dilation, Array<IndexExpr> padding, String layout, |
1086 | String out_layout, bool ceil_mode) { |
1087 | return MakeMaxPool<MaxPool1DAttrs>(data, pool_size, strides, dilation, padding, layout, |
1088 | out_layout, ceil_mode, "nn.max_pool1d" ); |
1089 | }); |
1090 | |
1091 | RELAY_REGISTER_OP("nn.max_pool1d" ) |
1092 | .describe(R"code(Max pooling operation for one dimensional data. |
1093 | |
1094 | - **data**: This depends on the `layout` parameter. Input is 3D array of shape |
1095 | (batch_size, channels, width) if `layout` is `NCW`. |
1096 | - **out**: This depends on the `layout` parameter. Output is 3D array of shape |
1097 | (batch_size, channels, , out_width) if `layout` is `NCW`. |
1098 | out_width is calculated as:: |
1099 | |
1100 | out_width = floor((width+padding[0]+padding[1]-pool_size[0])/strides[0])+1 |
1101 | |
1102 | where padding will be an expanded array based on number of values passed as:: |
1103 | one int : all sides same padding used. |
1104 | two int: padding width in the order of (left, right). |
1105 | |
1106 | When `ceil_mode` is `True`, ceil will be used instead of floor in this |
1107 | equation. |
1108 | |
1109 | )code" TVM_ADD_FILELINE) |
1110 | .set_attrs_type<MaxPool1DAttrs>() |
1111 | .set_num_inputs(1) |
1112 | .add_argument("data" , "Tensor" , "The input tensor." ) |
1113 | .set_support_level(2) |
1114 | .add_type_rel("MaxPool1D" , Pool1DRel<MaxPool1DAttrs>) |
1115 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , PoolInferCorrectLayout<MaxPool1DAttrs>) |
1116 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
1117 | .set_attr<FTVMCompute>("FTVMCompute" , Pool1DCompute<MaxPool1DAttrs, topi::nn::kMaxPool>); |
1118 | |
1119 | // AvgPool1D |
1120 | TVM_REGISTER_GLOBAL("relay.op.nn._make.avg_pool1d" ) |
1121 | .set_body_typed([](Expr data, Array<IndexExpr> pool_size, Array<IndexExpr> strides, |
1122 | Array<IndexExpr> dilation, Array<IndexExpr> padding, String layout, |
1123 | String out_layout, bool ceil_mode, bool count_include_pad) { |
1124 | return MakeAvgPool<AvgPool1DAttrs>(data, pool_size, strides, dilation, padding, layout, |
1125 | out_layout, ceil_mode, count_include_pad, "nn.avg_pool1d" ); |
1126 | }); |
1127 | |
1128 | RELAY_REGISTER_OP("nn.avg_pool1d" ) |
1129 | .describe(R"code( |
1130 | Average pooling operation for one dimensional data. |
1131 | |
1132 | - **data**: This depends on the `layout` parameter. Input is 3D array of shape |
1133 | (batch_size, channels, width) if `layout` is `NCW`. |
1134 | - **out**: This depends on the `layout` parameter. Output is 3D array of shape |
1135 | (batch_size, channels, out_width) if `layout` is `NCW`. |
1136 | out_width is calculated as:: |
1137 | |
1138 | out_width = floor((width+padding[0]+padding[1]-pool_size[0])/strides[0])+1 |
1139 | |
1140 | where padding will be an expanded array based on number of values passed as:: |
1141 | one int : all sides same padding used. |
1142 | two int: padding width in the order of (left, right). |
1143 | |
1144 | When `ceil_mode` is `True`, ceil will be used instead of floor in this |
1145 | equation. |
1146 | |
1147 | )code" TVM_ADD_FILELINE) |
1148 | .set_attrs_type<AvgPool1DAttrs>() |
1149 | .set_num_inputs(1) |
1150 | .add_argument("data" , "Tensor" , "The input tensor." ) |
1151 | .set_support_level(2) |
1152 | .add_type_rel("AvgPool1D" , Pool1DRel<AvgPool1DAttrs>) |
1153 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , PoolInferCorrectLayout<AvgPool1DAttrs>) |
1154 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
1155 | .set_attr<FTVMCompute>("FTVMCompute" , Pool1DCompute<AvgPool1DAttrs, topi::nn::kAvgPool>); |
1156 | |
1157 | // relay.nn.max_pool3d & relay.nn.avg_pool3d |
1158 | TVM_REGISTER_NODE_TYPE(MaxPool3DAttrs); |
1159 | TVM_REGISTER_NODE_TYPE(AvgPool3DAttrs); |
1160 | |
1161 | template <typename AttrType> |
1162 | bool Pool3DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, |
1163 | const TypeReporter& reporter) { |
1164 | ICHECK_EQ(types.size(), 2); |
1165 | const auto* data = types[0].as<TensorTypeNode>(); |
1166 | |
1167 | if (data == nullptr) return false; |
1168 | |
1169 | const auto dshape = data->shape; |
1170 | ICHECK_GE(dshape.size(), 3U) |
1171 | << "Pool3D only support input >= 3-D: input must have depth, height and width" ; |
1172 | const auto param = attrs.as<AttrType>(); |
1173 | ICHECK(param != nullptr); |
1174 | |
1175 | Layout layout(param->layout); |
1176 | Layout out_layout(param->out_layout); |
1177 | ICHECK(layout.Contains(LayoutAxis::Get('D')) && layout.Contains(LayoutAxis::Get('H')) && |
1178 | layout.Contains(LayoutAxis::Get('W')) && !layout.Contains(LayoutAxis::Get('d')) && |
1179 | !layout.Contains(LayoutAxis::Get('h')) && !layout.Contains(LayoutAxis::Get('w'))) |
1180 | << "Invalid layout " << layout |
1181 | << ". Pool3D layout must have D, H and W, which cannot be split" ; |
1182 | |
1183 | const auto didx = layout.IndexOf(LayoutAxis::Get('D')); |
1184 | const auto hidx = layout.IndexOf(LayoutAxis::Get('H')); |
1185 | const auto widx = layout.IndexOf(LayoutAxis::Get('W')); |
1186 | |
1187 | IndexExpr pad[3]; |
1188 | if (param->padding.size() == 1) { |
1189 | pad[0] = param->padding[0] * 2; |
1190 | pad[1] = param->padding[0] * 2; |
1191 | pad[2] = param->padding[0] * 2; |
1192 | } else if (param->padding.size() == 3) { |
1193 | // (front, top, left) |
1194 | pad[0] = param->padding[0] * 2; |
1195 | pad[1] = param->padding[1] * 2; |
1196 | pad[2] = param->padding[2] * 2; |
1197 | } else if (param->padding.size() == 6) { |
1198 | // (front, top, left, back, bottom, right) |
1199 | pad[0] = param->padding[0] + param->padding[3]; |
1200 | pad[1] = param->padding[1] + param->padding[4]; |
1201 | pad[2] = param->padding[2] + param->padding[5]; |
1202 | } else { |
1203 | return false; |
1204 | } |
1205 | |
1206 | std::vector<IndexExpr> oshape(dshape.begin(), dshape.end()); |
1207 | |
1208 | int idxes[3] = {didx, hidx, widx}; |
1209 | for (int i = 0; i < 3; i++) { |
1210 | int ii = idxes[i]; |
1211 | if (dshape[ii].as<tir::AnyNode>()) { |
1212 | oshape[ii] = dshape[ii]; |
1213 | } else { |
1214 | oshape[ii] = |
1215 | calculate_pool_dimension(dshape[ii], pad[i], param->pool_size[i], param->dilation[i], |
1216 | param->strides[i], param->ceil_mode); |
1217 | } |
1218 | } |
1219 | |
1220 | // assign output type |
1221 | reporter->Assign(types[1], TensorType(oshape, data->dtype)); |
1222 | return true; |
1223 | } |
1224 | |
1225 | template <typename AttrType, topi::nn::PoolType mode> |
1226 | Array<te::Tensor> Pool3DCompute(const Attrs& attrs, const Array<te::Tensor>& inputs, |
1227 | const Type& out_type) { |
1228 | static const Layout kNCDHW("NCDHW" ); |
1229 | const auto* param = attrs.as<AttrType>(); |
1230 | ICHECK(param != nullptr); |
1231 | auto pool_size = param->pool_size; |
1232 | auto strides = param->strides; |
1233 | auto dilation = param->dilation; |
1234 | auto padding = param->padding; |
1235 | auto ceil_mode = param->ceil_mode; |
1236 | Layout layout(param->layout); |
1237 | Layout out_layout(param->out_layout); |
1238 | |
1239 | ICHECK(tir::BijectiveLayout(layout, kNCDHW).defined()) |
1240 | << "max_pool3d currently only supports layouts that are convertible from NCDHW" ; |
1241 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('d')), -1) |
1242 | << "max_pool3d does not support input split on depth" ; |
1243 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('h')), -1) |
1244 | << "max_pool3d does not support input split on height" ; |
1245 | ICHECK_EQ(layout.IndexOf(LayoutAxis::Get('w')), -1) |
1246 | << "max_pool3d does not support input split on width" ; |
1247 | |
1248 | ICHECK(inputs[0].ndim() == 4U || inputs[0].ndim() == 5U || inputs[0].ndim() == 6U) |
1249 | << "Pool3D only support 5-D input (e.g., NCDHW)" |
1250 | << " or 6-D input (e.g. NCDHWc on for vector instructions)" |
1251 | << " or 7-D input (e.g. NCDHWnc for tensor accelerators)" ; |
1252 | |
1253 | if (param->padding.size() == 1) { |
1254 | padding.push_back(padding[0]); |
1255 | padding.push_back(padding[0]); |
1256 | padding.push_back(padding[0]); |
1257 | } else if (param->padding.size() == 3) { |
1258 | padding.push_back(padding[0]); |
1259 | padding.push_back(padding[1]); |
1260 | padding.push_back(padding[2]); |
1261 | } |
1262 | if (mode == topi::nn::kAvgPool) { |
1263 | bool count_include_pad = reinterpret_cast<const AvgPool3DAttrs*>(param)->count_include_pad; |
1264 | return Array<te::Tensor>{topi::nn::pool3d(inputs[0], pool_size, strides, dilation, padding, |
1265 | mode, ceil_mode, layout.name(), count_include_pad)}; |
1266 | } else { |
1267 | return Array<te::Tensor>{topi::nn::pool3d(inputs[0], pool_size, strides, dilation, padding, |
1268 | mode, ceil_mode, layout.name())}; |
1269 | } |
1270 | } |
1271 | |
1272 | TVM_REGISTER_GLOBAL("relay.op.nn._make.max_pool3d" ) |
1273 | .set_body_typed([](Expr data, Array<IndexExpr> pool_size, Array<IndexExpr> strides, |
1274 | Array<IndexExpr> dilation, Array<IndexExpr> padding, String layout, |
1275 | String out_layout, bool ceil_mode) { |
1276 | return MakeMaxPool<MaxPool3DAttrs>(data, pool_size, strides, dilation, padding, layout, |
1277 | out_layout, ceil_mode, "nn.max_pool3d" ); |
1278 | }); |
1279 | |
1280 | RELAY_REGISTER_OP("nn.max_pool3d" ) |
1281 | .describe(R"code(Max pooling operation for three dimensional data. |
1282 | |
1283 | - **data**: This depends on the `layout` parameter. Input is 5D array of shape |
1284 | (batch_size, channels, depth, height, width) if `layout` is `NCDHW`. |
1285 | - **out**: This depends on the `layout` parameter. Output is 5D array of shape |
1286 | (batch_size, channels, out_depth, out_height, out_width) if `layout` is `NCDHW`. |
1287 | out_depth, out_height and out_width are calculated as:: |
1288 | |
1289 | out_depth = floor((depth+padding[0]+padding[3]-pool_size[0])/strides[0])+1 |
1290 | out_height = floor((height+padding[1]+padding[4]-pool_size[1])/strides[1])+1 |
1291 | out_width = floor((width+padding[2]+padding[5]-pool_size[2])/strides[2])+1 |
1292 | |
1293 | where padding will be an expanded array based on number of values passed as:: |
1294 | one int : all sides same padding used. |
1295 | three int : front, bottom, right use same as back, top and left. |
1296 | six int: padding width in the order of (front, top, left, back, bottom, right). |
1297 | |
1298 | When `ceil_mode` is `True`, ceil will be used instead of floor in this |
1299 | equation. |
1300 | |
1301 | )code" TVM_ADD_FILELINE) |
1302 | .set_attrs_type<MaxPool3DAttrs>() |
1303 | .set_num_inputs(1) |
1304 | .add_argument("data" , "Tensor" , "The input tensor." ) |
1305 | .set_support_level(2) |
1306 | .add_type_rel("MaxPool3D" , Pool3DRel<MaxPool3DAttrs>) |
1307 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , PoolInferCorrectLayout<MaxPool3DAttrs>) |
1308 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
1309 | .set_attr<FTVMCompute>("FTVMCompute" , Pool3DCompute<MaxPool3DAttrs, topi::nn::kMaxPool>); |
1310 | |
1311 | // AvgPool3D |
1312 | TVM_REGISTER_GLOBAL("relay.op.nn._make.avg_pool3d" ) |
1313 | .set_body_typed([](Expr data, Array<IndexExpr> pool_size, Array<IndexExpr> strides, |
1314 | Array<IndexExpr> dilation, Array<IndexExpr> padding, String layout, |
1315 | String out_layout, bool ceil_mode, bool count_include_pad) { |
1316 | return MakeAvgPool<AvgPool3DAttrs>(data, pool_size, strides, dilation, padding, layout, |
1317 | out_layout, ceil_mode, count_include_pad, "nn.avg_pool3d" ); |
1318 | }); |
1319 | |
1320 | RELAY_REGISTER_OP("nn.avg_pool3d" ) |
1321 | .describe(R"code( |
1322 | Average pooling operation for three dimensional data. |
1323 | |
1324 | - **data**: This depends on the `layout` parameter. Input is 5D array of shape |
1325 | (batch_size, channels, depth, height, width) if `layout` is `NCDHW`. |
1326 | - **out**: This depends on the `layout` parameter. Output is 5D array of shape |
1327 | (batch_size, channels, out_depth, out_height, out_width) if `layout` is `NCDHW`. |
1328 | out_depth, out_height and out_width are calculated as:: |
1329 | |
1330 | out_depth = floor((depth+padding[0]+padding[3]-pool_size[0])/strides[0])+1 |
1331 | out_height = floor((height+padding[1]+padding[4]-pool_size[1])/strides[1])+1 |
1332 | out_width = floor((width+padding[2]+padding[5]-pool_size[2])/strides[2])+1 |
1333 | |
1334 | where padding will be an expanded array based on number of values passed as:: |
1335 | one int : all sides same padding used. |
1336 | three int : front, bottom, right use same as back, top and left. |
1337 | six int: padding width in the order of (front, top, left, back, bottom, right). |
1338 | |
1339 | When `ceil_mode` is `True`, ceil will be used instead of floor in this |
1340 | equation. |
1341 | |
1342 | )code" TVM_ADD_FILELINE) |
1343 | .set_attrs_type<AvgPool3DAttrs>() |
1344 | .set_num_inputs(1) |
1345 | .add_argument("data" , "Tensor" , "The input tensor." ) |
1346 | .set_support_level(2) |
1347 | .add_type_rel("AvgPool3D" , Pool3DRel<AvgPool3DAttrs>) |
1348 | .set_attr<FInferCorrectLayout>("FInferCorrectLayout" , PoolInferCorrectLayout<AvgPool3DAttrs>) |
1349 | .set_attr<TOpPattern>("TOpPattern" , kOutEWiseFusable) |
1350 | .set_attr<FTVMCompute>("FTVMCompute" , Pool3DCompute<AvgPool3DAttrs, topi::nn::kAvgPool>); |
1351 | |
1352 | } // namespace relay |
1353 | } // namespace tvm |
1354 | |