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
36namespace tvm {
37namespace relay {
38
39// relay.nn.max_pool2d & relay.nn.avg_pool2d
40TVM_REGISTER_NODE_TYPE(MaxPool2DAttrs);
41TVM_REGISTER_NODE_TYPE(AvgPool2DAttrs);
42
43template <typename T>
44InferCorrectLayoutOutput 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
66IndexExpr 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
80template <typename AttrType>
81bool 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
140template <typename AttrType, topi::nn::PoolType mode>
141Array<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
184TVM_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
192RELAY_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
223TVM_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
231RELAY_REGISTER_OP("nn.avg_pool2d")
232 .describe(R"code(
233Average 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
263TVM_REGISTER_NODE_TYPE(GlobalPool2DAttrs);
264
265bool 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
294template <topi::nn::PoolType mode>
295Array<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
314Expr 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
322TVM_REGISTER_GLOBAL("relay.op.nn._make.global_avg_pool2d").set_body_typed(MakeGlobalAvgPool2D);
323
324// GlobalAvgPool
325RELAY_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
344Expr 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
352TVM_REGISTER_GLOBAL("relay.op.nn._make.global_max_pool2d").set_body_typed(MakeGlobalMaxPool2D);
353
354RELAY_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
373TVM_REGISTER_NODE_TYPE(AdaptivePool1DAttrs);
374
375bool 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
409template <topi::nn::PoolType mode>
410Array<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
438Expr 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
448TVM_REGISTER_GLOBAL("relay.op.nn._make.adaptive_avg_pool1d").set_body_typed(MakeAdaptiveAvgPool1D);
449
450RELAY_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
474Expr 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
484TVM_REGISTER_GLOBAL("relay.op.nn._make.adaptive_max_pool1d").set_body_typed(MakeAdaptiveMaxPool1D);
485
486RELAY_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
510TVM_REGISTER_NODE_TYPE(AdaptivePool2DAttrs);
511
512bool 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
555template <topi::nn::PoolType mode>
556Array<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
592Expr 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
602TVM_REGISTER_GLOBAL("relay.op.nn._make.adaptive_avg_pool2d").set_body_typed(MakeAdaptiveAvgPool2D);
603
604RELAY_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
630Expr 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
640TVM_REGISTER_GLOBAL("relay.op.nn._make.adaptive_max_pool2d").set_body_typed(MakeAdaptiveMaxPool2D);
641
642RELAY_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
668TVM_REGISTER_NODE_TYPE(AdaptivePool3DAttrs);
669
670bool 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
720template <topi::nn::PoolType mode>
721Array<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
765Expr 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
775TVM_REGISTER_GLOBAL("relay.op.nn._make.adaptive_max_pool3d").set_body_typed(MakeAdaptiveMaxPool3D);
776
777RELAY_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
803Expr 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
813TVM_REGISTER_GLOBAL("relay.op.nn._make.adaptive_avg_pool3d").set_body_typed(MakeAdaptiveAvgPool3D);
814
815RELAY_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
838bool 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
850template <typename AttrType, topi::nn::PoolType mode>
851Array<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
898Expr 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
912TVM_REGISTER_GLOBAL("relay.op.nn._make.max_pool2d_grad").set_body_typed(MakeMaxPool2DGrad);
913
914RELAY_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
947Expr 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
962TVM_REGISTER_GLOBAL("relay.op.nn._make.avg_pool2d_grad").set_body_typed(MakeAvgPool2DGrad);
963
964RELAY_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
997TVM_REGISTER_NODE_TYPE(MaxPool1DAttrs);
998TVM_REGISTER_NODE_TYPE(AvgPool1DAttrs);
999
1000template <typename AttrType>
1001bool 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
1045template <typename AttrType, topi::nn::PoolType mode>
1046Array<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
1083TVM_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
1091RELAY_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
1120TVM_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
1128RELAY_REGISTER_OP("nn.avg_pool1d")
1129 .describe(R"code(
1130Average 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
1158TVM_REGISTER_NODE_TYPE(MaxPool3DAttrs);
1159TVM_REGISTER_NODE_TYPE(AvgPool3DAttrs);
1160
1161template <typename AttrType>
1162bool 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
1225template <typename AttrType, topi::nn::PoolType mode>
1226Array<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
1272TVM_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
1280RELAY_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
1312TVM_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
1320RELAY_REGISTER_OP("nn.avg_pool3d")
1321 .describe(R"code(
1322Average 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