[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op

[GitBox <gi...@apache.org>]

masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r355131252
 
 

 ##########
 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:
   This is a minor function, you can remove find_height_width function and use this one for 2d

----------------------------------------------------------------
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:
users@infra.apache.org


With regards,
Apache Git Services