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