optima2005 commented on a change in pull request #4478: [TOPI] implement pool3d
op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r355154973
##########
File path: topi/include/topi/nn/pooling.h
##########
@@ -591,6 +593,211 @@ inline Tensor global_pool(const Tensor& x,
return adaptive_pool(x, Array<Expr>{1, 1}, pool_type, layout);
}
+/*!
+* \brief Perform pooling on N-dimension of data.
+*
+* \param x The input tensor
+* \param kernel_size Vector of N ints
+* \param stride_size Vector of N ints
+* \param padding_size Vector of N*2 ints [head_pad_d1, head_pad_d2, ...,
+* head_pad_dN, tail_pad_d1, tail_pad_d2, ..., tail_pad_dN]
+* \param pool_type The type of pooling operator
+* \param ceil_mode Whether to use ceil when calculating the output size
+* \param axis Vector of indices for the N dimensions
+* \param count_include_pad Whether include padding in the calculation
+*
+* \return The output tensor in same layout order
+*/
+inline Tensor pool_impl_nd(const Tensor& x,
+ const Array<Expr>& kernel_size,
+ const Array<Expr>& stride_size,
+ const Array<Expr>& padding_size,
+ PoolType pool_type,
+ bool ceil_mode,
+ const std::vector<int>& axis,
+ bool count_include_pad) {
+ int k_size = kernel_size.size();
+ int x_size = x->shape.size();
+ CHECK_EQ(stride_size.size(), k_size) << "Pooling stride_size must have same
elements as kernel";
+ CHECK_EQ(padding_size.size(), k_size * 2) << "Pooling padding_size must has
double elements of"
+ " kernel";
+ CHECK_EQ(axis.size(), k_size) << "axis must have same elements as kernel";
+
+ Array<IterVar> daxis;
+ std::vector<Expr> kernel(k_size);
+ std::vector<Expr> stride(k_size);
+ std::vector<Expr> pad_head(k_size);
+ std::vector<Expr> pad_tail(k_size);
+ Array<Expr> pad_before(std::vector<Expr>(x_size, 0));
+ Array<Expr> pad_after(std::vector<Expr>(x_size, 0));
+ Array<Expr> out_shape = x->shape;
+
+ bool do_pad = false;
+ for (int i = 0; i < k_size; i++) {
+ int ii = axis[i];
+ kernel[i] = cast(Int(32), kernel_size[i]);
+ stride[i] = cast(Int(32), stride_size[i]);
+ pad_head[i] = cast(Int(32), padding_size[i]);
+ pad_tail[i] = cast(Int(32), padding_size[i + k_size]);
+ const int64_t *padding0 = as_const_int(pad_head[i]);
+ const int64_t *padding1 = as_const_int(pad_tail[i]);
+ do_pad = (do_pad) ? do_pad : ((padding0 && *padding0) || (padding1 &&
*padding1));
+
+ if (ceil_mode) {
+ // Additional padding to ensure we do ceil instead of floor when
+ // dividing by stride.
+ pad_tail[i] += stride[i] - 1;
+ }
+
+ daxis.push_back(tvm::reduce_axis(Range(0, kernel[i])));
+
+ pad_before.Set(ii, pad_head[i]);
+ pad_after.Set(ii, pad_tail[i]);
+
+ auto out_dim = tvm::ir::Simplify(
+ indexdiv(x->shape[ii] - kernel[i] + pad_head[i] + pad_tail[i],
stride[i]) + 1);
+
+ out_shape.Set(ii, out_dim);
+ }
+
+ if (pool_type == kMaxPool) {
+ auto temp = do_pad ? pad(x, pad_before, pad_after, x->dtype.min(),
"pad_temp") : x;
+ return tvm::compute(out_shape, [&](const Array<Var>& output) {
+ Array<Expr> indices;
+ for (const Var& var : output) indices.push_back(var);
+
+ for (int i = 0; i < k_size; i++) {
+ int ii = axis[i];
+ indices.Set(ii, output[ii] * stride[i] + daxis[i]);
+ }
+
+ return tvm::max(temp(indices), daxis);
+ }, "tensor", "pool_max");
+ } else if (pool_type == kAvgPool) {
+ // Pad the inputs
+ auto temp = do_pad ? pad(x, pad_before, pad_after, 0, "pad_temp") : x;
+
+ // TVM compute for summing the pooling window.
+ auto pool_sum = tvm::compute(out_shape,
+ [&](const Array<Var>& output) {
+ Array<Expr> indices;
+ for (const Var& var : output) indices.push_back(var);
+
+ for (int i = 0; i < k_size; i++) {
+ int ii = axis[i];
+ indices.Set(ii, output[ii] * stride[i] + daxis[i]);
+ }
+ return tvm::sum(temp(indices), daxis);
+ }, "tensor", "pool_sum");
+
+ // TVM compute for dividing the reduced window sum by kernel size.
+ return tvm::compute(out_shape,
+ [&](const Array<Var>& output) {
+ Array<Expr> indices;
+ for (const Var& var : output) indices.push_back(var);
+ if (count_include_pad) {
+ auto kernel_size = make_const(Int(32), 1);
+ for (int i = 0; i < k_size; i++) {
+ kernel_size *= kernel[i];
+ }
+ return div(pool_sum(indices), kernel_size);
+ } else {
+ std::vector<Expr> start(k_size);
+ std::vector<Expr> end(k_size);
+ auto kernel_size = make_const(Int(32), 1);
+ for (int i = 0; i < k_size; i++) {
+ int ii = axis[i];
+ start[i] = output[ii] * stride[i] - pad_head[i];
+ end[i] = ir::Min::make(start[i] + kernel[i], x->shape[ii]);
+ start[i] = ir::Max::make(start[i], make_const(Int(32), 0));
+ kernel_size *= (end[i] - start[i]);
+ }
+
+ Expr divide_factor = ir::Max::make(kernel_size, make_const(Int(32),
1));
+ return div(pool_sum(indices), divide_factor);
+ }
+ }, "tensor", kElementWise);
+ } else {
+ LOG(ERROR) << "Unrecognized pool_type: " << pool_type;
+ return x;
+ }
+}
+
+inline bool find_depth_height_width(const std::string& layout,
Review comment:
The number of return variables are different. The extra logics would be
needed to handle this difference. I am wondering whether it is worth to do
that.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
