[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r356942674
##
File path: topi/include/topi/nn/pooling.h
##
@@ -591,6 +593,211 @@ inline Tensor global_pool(const Tensor& x,
return adaptive_pool(x, Array{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& kernel_size,
+ const Array& stride_size,
+ const Array& padding_size,
+ PoolType pool_type,
+ bool ceil_mode,
+ const std::vector& 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 daxis;
+ std::vector kernel(k_size);
+ std::vector stride(k_size);
+ std::vector pad_head(k_size);
+ std::vector pad_tail(k_size);
+ Array pad_before(std::vector(x_size, 0));
+ Array pad_after(std::vector(x_size, 0));
+ Array 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& output) {
+ Array 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& output) {
+ Array 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& output) {
+ Array 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 start(k_size);
+std::vector 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
[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r356942674
##
File path: topi/include/topi/nn/pooling.h
##
@@ -591,6 +593,211 @@ inline Tensor global_pool(const Tensor& x,
return adaptive_pool(x, Array{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& kernel_size,
+ const Array& stride_size,
+ const Array& padding_size,
+ PoolType pool_type,
+ bool ceil_mode,
+ const std::vector& 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 daxis;
+ std::vector kernel(k_size);
+ std::vector stride(k_size);
+ std::vector pad_head(k_size);
+ std::vector pad_tail(k_size);
+ Array pad_before(std::vector(x_size, 0));
+ Array pad_after(std::vector(x_size, 0));
+ Array 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& output) {
+ Array 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& output) {
+ Array 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& output) {
+ Array 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 start(k_size);
+std::vector 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
[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r356942674
##
File path: topi/include/topi/nn/pooling.h
##
@@ -591,6 +593,211 @@ inline Tensor global_pool(const Tensor& x,
return adaptive_pool(x, Array{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& kernel_size,
+ const Array& stride_size,
+ const Array& padding_size,
+ PoolType pool_type,
+ bool ceil_mode,
+ const std::vector& 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 daxis;
+ std::vector kernel(k_size);
+ std::vector stride(k_size);
+ std::vector pad_head(k_size);
+ std::vector pad_tail(k_size);
+ Array pad_before(std::vector(x_size, 0));
+ Array pad_after(std::vector(x_size, 0));
+ Array 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& output) {
+ Array 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& output) {
+ Array 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& output) {
+ Array 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 start(k_size);
+std::vector 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
[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r356740860
##
File path: topi/include/topi/nn/pooling.h
##
@@ -591,6 +593,211 @@ inline Tensor global_pool(const Tensor& x,
return adaptive_pool(x, Array{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& kernel_size,
+ const Array& stride_size,
+ const Array& padding_size,
+ PoolType pool_type,
+ bool ceil_mode,
+ const std::vector& 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 daxis;
+ std::vector kernel(k_size);
+ std::vector stride(k_size);
+ std::vector pad_head(k_size);
+ std::vector pad_tail(k_size);
+ Array pad_before(std::vector(x_size, 0));
+ Array pad_after(std::vector(x_size, 0));
+ Array 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& output) {
+ Array 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& output) {
+ Array 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& output) {
+ Array 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 start(k_size);
+std::vector 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
[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r356739029
##
File path: src/relay/op/nn/pooling.cc
##
@@ -720,5 +720,238 @@ RELAY_REGISTER_OP("nn.avg_pool2d_grad")
.set_attr("FTVMCompute", Pool2DGradCompute);
+// relay.nn.max_pool3d & relay.nn.avg_pool3d
+TVM_REGISTER_NODE_TYPE(MaxPool3DAttrs);
+TVM_REGISTER_NODE_TYPE(AvgPool3DAttrs);
+
+template
+bool Pool3DRel(const Array& types,
+ int num_inputs,
+ const Attrs& attrs,
+ const TypeReporter& reporter) {
+ CHECK_EQ(types.size(), 2);
+ const auto* data = types[0].as();
+
+ if (data == nullptr) return false;
+
+ const auto dshape = data->shape;
+ CHECK_GE(dshape.size(), 3U)
+ << "Pool3D only support input >= 3-D: input must have depth, height and
width";
+ const auto param = attrs.as();
+ CHECK(param != nullptr);
+
+ Layout layout(param->layout);
+ CHECK(layout.Contains(LayoutAxis::Get('D')) &&
layout.Contains(LayoutAxis::Get('H')) &&
+layout.Contains(LayoutAxis::Get('W')) &&
!layout.Contains(LayoutAxis::Get('d')) &&
+!layout.Contains(LayoutAxis::Get('h')) &&
!layout.Contains(LayoutAxis::Get('w')))
+<< "Invalid layout " << layout
+<< ". Pool3D layout must have D, H and W, which cannot be split";
+
+ const auto didx = layout.IndexOf(LayoutAxis::Get('D'));
+ const auto hidx = layout.IndexOf(LayoutAxis::Get('H'));
+ const auto widx = layout.IndexOf(LayoutAxis::Get('W'));
+
+ IndexExpr pad_d, pad_h, pad_w;
+ if (param->padding.size() == 1) {
+pad_d = param->padding[0] * 2;
+pad_h = param->padding[0] * 2;
+pad_w = param->padding[0] * 2;
+ } else if (param->padding.size() == 3) {
+// (front, top, left)
+pad_d = param->padding[0] * 2;
+pad_h = param->padding[1] * 2;
+pad_w = param->padding[2] * 2;
+ } else if (param->padding.size() == 6) {
+// (front, top, left, back, bottom, right)
+pad_d = param->padding[0] + param->padding[3];
+pad_h = param->padding[1] + param->padding[4];
+pad_w = param->padding[2] + param->padding[5];
+ } else {
+return false;
+ }
+
+ std::vector oshape;
+ for (const auto& e : dshape) {
+oshape.push_back(e);
+ }
+
+ std::vector idxes = {didx, hidx, widx};
+ for (int i = 0; i < 3; i++) {
+int ii = idxes[i];
+if (dshape[ii].as()) {
+ oshape[ii] = dshape[ii];
+} else {
+ if (param->ceil_mode) {
+oshape[ii] = ((dshape[ii] + pad_d - param->pool_size[i] +
+ param->strides[i] - 1) / param->strides[i]) + 1;
+ } else {
+oshape[ii] = ((dshape[ii] + pad_d - param->pool_size[i]) /
param->strides[i]) + 1;
+ }
+}
+ }
+
+ // assign output type
+ reporter->Assign(types[1], TensorTypeNode::make(oshape, data->dtype));
+ return true;
+}
+
+// MaxPool3D
+Expr MakeMaxPool3D(Expr data,
+ Array pool_size,
+ Array strides,
+ Array padding,
+ std::string layout,
+ bool ceil_mode) {
+ auto attrs = make_node();
+ attrs->pool_size = std::move(pool_size);
+ attrs->strides = std::move(strides);
+ attrs->padding = std::move(padding);
+ attrs->layout = std::move(layout);
+ attrs->ceil_mode = ceil_mode;
+ static const Op& op = Op::Get("nn.max_pool3d");
+ return CallNode::make(op, {data}, Attrs(attrs), {});
+}
+
+template
+Array Pool3DCompute(const Attrs& attrs,
+const Array& inputs,
+const Type& out_type,
+const Target& target) {
+ static const Layout kNCDHW("NCDHW");
+ const auto* param = attrs.as();
+ CHECK(param != nullptr);
+ auto pool_size = param->pool_size;
+ auto strides = param->strides;
+ auto padding = param->padding;
+ auto ceil_mode = param->ceil_mode;
+ Layout layout(param->layout);
+
+ CHECK(BijectiveLayoutNode::make(layout, kNCDHW).defined())
+ << "max_pool3d currently only supports layouts that are convertible from
NCDHW";
+ CHECK_EQ(layout.IndexOf(LayoutAxis::Get('d')), -1)
+ << "max_pool3d does not support input split on depth";
+ CHECK_EQ(layout.IndexOf(LayoutAxis::Get('h')), -1)
+ << "max_pool3d does not support input split on height";
+ CHECK_EQ(layout.IndexOf(LayoutAxis::Get('w')), -1)
+ << "max_pool3d does not support input split on width";
+
+ CHECK(inputs[0].ndim() == 4U ||
+inputs[0].ndim() == 5U ||
+inputs[0].ndim() == 6U)
+ << "Pool3D only support 5-D input (e.g., NCDHW)"
+ << " or 6-D input (e.g. NCDHWc on for vector instructions)"
+ << " or 7-D input (e.g. NCDHWnc for tensor accelerators)";
+
+ if (param->padding.size() == 1) {
+padding.push_back(padding[0]);
+padding.push_back(padding[0]);
+padding.push_back(padding[0]);
+ } else if (param->padding.size() == 3) {
+padding.push_back(padding[0]);
+padding.push_back(padding[1]);
[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r356738841
##
File path: src/relay/op/nn/pooling.cc
##
@@ -720,5 +720,238 @@ RELAY_REGISTER_OP("nn.avg_pool2d_grad")
.set_attr("FTVMCompute", Pool2DGradCompute);
+// relay.nn.max_pool3d & relay.nn.avg_pool3d
+TVM_REGISTER_NODE_TYPE(MaxPool3DAttrs);
+TVM_REGISTER_NODE_TYPE(AvgPool3DAttrs);
+
+template
+bool Pool3DRel(const Array& types,
+ int num_inputs,
+ const Attrs& attrs,
+ const TypeReporter& reporter) {
+ CHECK_EQ(types.size(), 2);
+ const auto* data = types[0].as();
+
+ if (data == nullptr) return false;
+
+ const auto dshape = data->shape;
+ CHECK_GE(dshape.size(), 3U)
+ << "Pool3D only support input >= 3-D: input must have depth, height and
width";
+ const auto param = attrs.as();
+ CHECK(param != nullptr);
+
+ Layout layout(param->layout);
+ CHECK(layout.Contains(LayoutAxis::Get('D')) &&
layout.Contains(LayoutAxis::Get('H')) &&
+layout.Contains(LayoutAxis::Get('W')) &&
!layout.Contains(LayoutAxis::Get('d')) &&
+!layout.Contains(LayoutAxis::Get('h')) &&
!layout.Contains(LayoutAxis::Get('w')))
+<< "Invalid layout " << layout
+<< ". Pool3D layout must have D, H and W, which cannot be split";
+
+ const auto didx = layout.IndexOf(LayoutAxis::Get('D'));
+ const auto hidx = layout.IndexOf(LayoutAxis::Get('H'));
+ const auto widx = layout.IndexOf(LayoutAxis::Get('W'));
+
+ IndexExpr pad_d, pad_h, pad_w;
+ if (param->padding.size() == 1) {
+pad_d = param->padding[0] * 2;
+pad_h = param->padding[0] * 2;
+pad_w = param->padding[0] * 2;
+ } else if (param->padding.size() == 3) {
+// (front, top, left)
+pad_d = param->padding[0] * 2;
+pad_h = param->padding[1] * 2;
+pad_w = param->padding[2] * 2;
+ } else if (param->padding.size() == 6) {
+// (front, top, left, back, bottom, right)
+pad_d = param->padding[0] + param->padding[3];
+pad_h = param->padding[1] + param->padding[4];
+pad_w = param->padding[2] + param->padding[5];
+ } else {
+return false;
+ }
+
+ std::vector oshape;
+ for (const auto& e : dshape) {
+oshape.push_back(e);
+ }
+
+ std::vector idxes = {didx, hidx, widx};
+ for (int i = 0; i < 3; i++) {
+int ii = idxes[i];
+if (dshape[ii].as()) {
+ oshape[ii] = dshape[ii];
+} else {
+ if (param->ceil_mode) {
+oshape[ii] = ((dshape[ii] + pad_d - param->pool_size[i] +
+ param->strides[i] - 1) / param->strides[i]) + 1;
+ } else {
+oshape[ii] = ((dshape[ii] + pad_d - param->pool_size[i]) /
param->strides[i]) + 1;
+ }
+}
+ }
+
+ // assign output type
+ reporter->Assign(types[1], TensorTypeNode::make(oshape, data->dtype));
+ return true;
+}
+
+// MaxPool3D
+Expr MakeMaxPool3D(Expr data,
+ Array pool_size,
+ Array strides,
+ Array padding,
+ std::string layout,
+ bool ceil_mode) {
+ auto attrs = make_node();
+ attrs->pool_size = std::move(pool_size);
+ attrs->strides = std::move(strides);
+ attrs->padding = std::move(padding);
+ attrs->layout = std::move(layout);
+ attrs->ceil_mode = ceil_mode;
+ static const Op& op = Op::Get("nn.max_pool3d");
+ return CallNode::make(op, {data}, Attrs(attrs), {});
+}
+
+template
+Array Pool3DCompute(const Attrs& attrs,
+const Array& inputs,
+const Type& out_type,
+const Target& target) {
+ static const Layout kNCDHW("NCDHW");
+ const auto* param = attrs.as();
+ CHECK(param != nullptr);
+ auto pool_size = param->pool_size;
+ auto strides = param->strides;
+ auto padding = param->padding;
+ auto ceil_mode = param->ceil_mode;
+ Layout layout(param->layout);
+
+ CHECK(BijectiveLayoutNode::make(layout, kNCDHW).defined())
+ << "max_pool3d currently only supports layouts that are convertible from
NCDHW";
+ CHECK_EQ(layout.IndexOf(LayoutAxis::Get('d')), -1)
+ << "max_pool3d does not support input split on depth";
+ CHECK_EQ(layout.IndexOf(LayoutAxis::Get('h')), -1)
+ << "max_pool3d does not support input split on height";
+ CHECK_EQ(layout.IndexOf(LayoutAxis::Get('w')), -1)
+ << "max_pool3d does not support input split on width";
+
+ CHECK(inputs[0].ndim() == 4U ||
+inputs[0].ndim() == 5U ||
+inputs[0].ndim() == 6U)
+ << "Pool3D only support 5-D input (e.g., NCDHW)"
+ << " or 6-D input (e.g. NCDHWc on for vector instructions)"
+ << " or 7-D input (e.g. NCDHWnc for tensor accelerators)";
+
+ if (param->padding.size() == 1) {
+padding.push_back(padding[0]);
+padding.push_back(padding[0]);
+padding.push_back(padding[0]);
+ } else if (param->padding.size() == 3) {
+padding.push_back(padding[0]);
+padding.push_back(padding[1]);
[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r356737854
##
File path: src/relay/op/nn/pooling.cc
##
@@ -720,5 +720,238 @@ RELAY_REGISTER_OP("nn.avg_pool2d_grad")
.set_attr("FTVMCompute", Pool2DGradCompute);
+// relay.nn.max_pool3d & relay.nn.avg_pool3d
+TVM_REGISTER_NODE_TYPE(MaxPool3DAttrs);
+TVM_REGISTER_NODE_TYPE(AvgPool3DAttrs);
+
+template
+bool Pool3DRel(const Array& types,
+ int num_inputs,
+ const Attrs& attrs,
+ const TypeReporter& reporter) {
+ CHECK_EQ(types.size(), 2);
+ const auto* data = types[0].as();
+
+ if (data == nullptr) return false;
+
+ const auto dshape = data->shape;
+ CHECK_GE(dshape.size(), 3U)
+ << "Pool3D only support input >= 3-D: input must have depth, height and
width";
+ const auto param = attrs.as();
+ CHECK(param != nullptr);
+
+ Layout layout(param->layout);
+ CHECK(layout.Contains(LayoutAxis::Get('D')) &&
layout.Contains(LayoutAxis::Get('H')) &&
+layout.Contains(LayoutAxis::Get('W')) &&
!layout.Contains(LayoutAxis::Get('d')) &&
+!layout.Contains(LayoutAxis::Get('h')) &&
!layout.Contains(LayoutAxis::Get('w')))
+<< "Invalid layout " << layout
+<< ". Pool3D layout must have D, H and W, which cannot be split";
+
+ const auto didx = layout.IndexOf(LayoutAxis::Get('D'));
+ const auto hidx = layout.IndexOf(LayoutAxis::Get('H'));
+ const auto widx = layout.IndexOf(LayoutAxis::Get('W'));
+
+ IndexExpr pad_d, pad_h, pad_w;
+ if (param->padding.size() == 1) {
+pad_d = param->padding[0] * 2;
+pad_h = param->padding[0] * 2;
+pad_w = param->padding[0] * 2;
+ } else if (param->padding.size() == 3) {
+// (front, top, left)
+pad_d = param->padding[0] * 2;
+pad_h = param->padding[1] * 2;
+pad_w = param->padding[2] * 2;
+ } else if (param->padding.size() == 6) {
+// (front, top, left, back, bottom, right)
+pad_d = param->padding[0] + param->padding[3];
+pad_h = param->padding[1] + param->padding[4];
+pad_w = param->padding[2] + param->padding[5];
+ } else {
+return false;
+ }
+
+ std::vector oshape;
+ for (const auto& e : dshape) {
+oshape.push_back(e);
+ }
+
+ std::vector idxes = {didx, hidx, widx};
+ for (int i = 0; i < 3; i++) {
+int ii = idxes[i];
+if (dshape[ii].as()) {
+ oshape[ii] = dshape[ii];
+} else {
+ if (param->ceil_mode) {
+oshape[ii] = ((dshape[ii] + pad_d - param->pool_size[i] +
+ param->strides[i] - 1) / param->strides[i]) + 1;
+ } else {
+oshape[ii] = ((dshape[ii] + pad_d - param->pool_size[i]) /
param->strides[i]) + 1;
+ }
+}
+ }
+
+ // assign output type
+ reporter->Assign(types[1], TensorTypeNode::make(oshape, data->dtype));
+ return true;
+}
+
+// MaxPool3D
+Expr MakeMaxPool3D(Expr data,
+ Array pool_size,
+ Array strides,
+ Array padding,
+ std::string layout,
+ bool ceil_mode) {
+ auto attrs = make_node();
+ attrs->pool_size = std::move(pool_size);
+ attrs->strides = std::move(strides);
+ attrs->padding = std::move(padding);
+ attrs->layout = std::move(layout);
+ attrs->ceil_mode = ceil_mode;
+ static const Op& op = Op::Get("nn.max_pool3d");
+ return CallNode::make(op, {data}, Attrs(attrs), {});
+}
+
+template
+Array Pool3DCompute(const Attrs& attrs,
+const Array& inputs,
+const Type& out_type,
+const Target& target) {
+ static const Layout kNCDHW("NCDHW");
+ const auto* param = attrs.as();
+ CHECK(param != nullptr);
+ auto pool_size = param->pool_size;
+ auto strides = param->strides;
+ auto padding = param->padding;
+ auto ceil_mode = param->ceil_mode;
+ Layout layout(param->layout);
+
+ CHECK(BijectiveLayoutNode::make(layout, kNCDHW).defined())
+ << "max_pool3d currently only supports layouts that are convertible from
NCDHW";
+ CHECK_EQ(layout.IndexOf(LayoutAxis::Get('d')), -1)
+ << "max_pool3d does not support input split on depth";
+ CHECK_EQ(layout.IndexOf(LayoutAxis::Get('h')), -1)
+ << "max_pool3d does not support input split on height";
+ CHECK_EQ(layout.IndexOf(LayoutAxis::Get('w')), -1)
+ << "max_pool3d does not support input split on width";
+
+ CHECK(inputs[0].ndim() == 4U ||
+inputs[0].ndim() == 5U ||
+inputs[0].ndim() == 6U)
+ << "Pool3D only support 5-D input (e.g., NCDHW)"
+ << " or 6-D input (e.g. NCDHWc on for vector instructions)"
+ << " or 7-D input (e.g. NCDHWnc for tensor accelerators)";
+
+ if (param->padding.size() == 1) {
+padding.push_back(padding[0]);
+padding.push_back(padding[0]);
+padding.push_back(padding[0]);
+ } else if (param->padding.size() == 3) {
+padding.push_back(padding[0]);
+padding.push_back(padding[1]);
[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r356737854
##
File path: src/relay/op/nn/pooling.cc
##
@@ -720,5 +720,238 @@ RELAY_REGISTER_OP("nn.avg_pool2d_grad")
.set_attr("FTVMCompute", Pool2DGradCompute);
+// relay.nn.max_pool3d & relay.nn.avg_pool3d
+TVM_REGISTER_NODE_TYPE(MaxPool3DAttrs);
+TVM_REGISTER_NODE_TYPE(AvgPool3DAttrs);
+
+template
+bool Pool3DRel(const Array& types,
+ int num_inputs,
+ const Attrs& attrs,
+ const TypeReporter& reporter) {
+ CHECK_EQ(types.size(), 2);
+ const auto* data = types[0].as();
+
+ if (data == nullptr) return false;
+
+ const auto dshape = data->shape;
+ CHECK_GE(dshape.size(), 3U)
+ << "Pool3D only support input >= 3-D: input must have depth, height and
width";
+ const auto param = attrs.as();
+ CHECK(param != nullptr);
+
+ Layout layout(param->layout);
+ CHECK(layout.Contains(LayoutAxis::Get('D')) &&
layout.Contains(LayoutAxis::Get('H')) &&
+layout.Contains(LayoutAxis::Get('W')) &&
!layout.Contains(LayoutAxis::Get('d')) &&
+!layout.Contains(LayoutAxis::Get('h')) &&
!layout.Contains(LayoutAxis::Get('w')))
+<< "Invalid layout " << layout
+<< ". Pool3D layout must have D, H and W, which cannot be split";
+
+ const auto didx = layout.IndexOf(LayoutAxis::Get('D'));
+ const auto hidx = layout.IndexOf(LayoutAxis::Get('H'));
+ const auto widx = layout.IndexOf(LayoutAxis::Get('W'));
+
+ IndexExpr pad_d, pad_h, pad_w;
+ if (param->padding.size() == 1) {
+pad_d = param->padding[0] * 2;
+pad_h = param->padding[0] * 2;
+pad_w = param->padding[0] * 2;
+ } else if (param->padding.size() == 3) {
+// (front, top, left)
+pad_d = param->padding[0] * 2;
+pad_h = param->padding[1] * 2;
+pad_w = param->padding[2] * 2;
+ } else if (param->padding.size() == 6) {
+// (front, top, left, back, bottom, right)
+pad_d = param->padding[0] + param->padding[3];
+pad_h = param->padding[1] + param->padding[4];
+pad_w = param->padding[2] + param->padding[5];
+ } else {
+return false;
+ }
+
+ std::vector oshape;
+ for (const auto& e : dshape) {
+oshape.push_back(e);
+ }
+
+ std::vector idxes = {didx, hidx, widx};
+ for (int i = 0; i < 3; i++) {
+int ii = idxes[i];
+if (dshape[ii].as()) {
+ oshape[ii] = dshape[ii];
+} else {
+ if (param->ceil_mode) {
+oshape[ii] = ((dshape[ii] + pad_d - param->pool_size[i] +
+ param->strides[i] - 1) / param->strides[i]) + 1;
+ } else {
+oshape[ii] = ((dshape[ii] + pad_d - param->pool_size[i]) /
param->strides[i]) + 1;
+ }
+}
+ }
+
+ // assign output type
+ reporter->Assign(types[1], TensorTypeNode::make(oshape, data->dtype));
+ return true;
+}
+
+// MaxPool3D
+Expr MakeMaxPool3D(Expr data,
+ Array pool_size,
+ Array strides,
+ Array padding,
+ std::string layout,
+ bool ceil_mode) {
+ auto attrs = make_node();
+ attrs->pool_size = std::move(pool_size);
+ attrs->strides = std::move(strides);
+ attrs->padding = std::move(padding);
+ attrs->layout = std::move(layout);
+ attrs->ceil_mode = ceil_mode;
+ static const Op& op = Op::Get("nn.max_pool3d");
+ return CallNode::make(op, {data}, Attrs(attrs), {});
+}
+
+template
+Array Pool3DCompute(const Attrs& attrs,
+const Array& inputs,
+const Type& out_type,
+const Target& target) {
+ static const Layout kNCDHW("NCDHW");
+ const auto* param = attrs.as();
+ CHECK(param != nullptr);
+ auto pool_size = param->pool_size;
+ auto strides = param->strides;
+ auto padding = param->padding;
+ auto ceil_mode = param->ceil_mode;
+ Layout layout(param->layout);
+
+ CHECK(BijectiveLayoutNode::make(layout, kNCDHW).defined())
+ << "max_pool3d currently only supports layouts that are convertible from
NCDHW";
+ CHECK_EQ(layout.IndexOf(LayoutAxis::Get('d')), -1)
+ << "max_pool3d does not support input split on depth";
+ CHECK_EQ(layout.IndexOf(LayoutAxis::Get('h')), -1)
+ << "max_pool3d does not support input split on height";
+ CHECK_EQ(layout.IndexOf(LayoutAxis::Get('w')), -1)
+ << "max_pool3d does not support input split on width";
+
+ CHECK(inputs[0].ndim() == 4U ||
+inputs[0].ndim() == 5U ||
+inputs[0].ndim() == 6U)
+ << "Pool3D only support 5-D input (e.g., NCDHW)"
+ << " or 6-D input (e.g. NCDHWc on for vector instructions)"
+ << " or 7-D input (e.g. NCDHWnc for tensor accelerators)";
+
+ if (param->padding.size() == 1) {
+padding.push_back(padding[0]);
+padding.push_back(padding[0]);
+padding.push_back(padding[0]);
+ } else if (param->padding.size() == 3) {
+padding.push_back(padding[0]);
+padding.push_back(padding[1]);
[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r356735856
##
File path: src/relay/op/nn/pooling.cc
##
@@ -720,5 +720,238 @@ RELAY_REGISTER_OP("nn.avg_pool2d_grad")
.set_attr("FTVMCompute", Pool2DGradCompute);
+// relay.nn.max_pool3d & relay.nn.avg_pool3d
+TVM_REGISTER_NODE_TYPE(MaxPool3DAttrs);
+TVM_REGISTER_NODE_TYPE(AvgPool3DAttrs);
+
+template
+bool Pool3DRel(const Array& types,
+ int num_inputs,
+ const Attrs& attrs,
+ const TypeReporter& reporter) {
+ CHECK_EQ(types.size(), 2);
+ const auto* data = types[0].as();
+
+ if (data == nullptr) return false;
+
+ const auto dshape = data->shape;
+ CHECK_GE(dshape.size(), 3U)
+ << "Pool3D only support input >= 3-D: input must have depth, height and
width";
+ const auto param = attrs.as();
+ CHECK(param != nullptr);
+
+ Layout layout(param->layout);
+ CHECK(layout.Contains(LayoutAxis::Get('D')) &&
layout.Contains(LayoutAxis::Get('H')) &&
+layout.Contains(LayoutAxis::Get('W')) &&
!layout.Contains(LayoutAxis::Get('d')) &&
+!layout.Contains(LayoutAxis::Get('h')) &&
!layout.Contains(LayoutAxis::Get('w')))
+<< "Invalid layout " << layout
+<< ". Pool3D layout must have D, H and W, which cannot be split";
+
+ const auto didx = layout.IndexOf(LayoutAxis::Get('D'));
+ const auto hidx = layout.IndexOf(LayoutAxis::Get('H'));
+ const auto widx = layout.IndexOf(LayoutAxis::Get('W'));
+
+ IndexExpr pad_d, pad_h, pad_w;
+ if (param->padding.size() == 1) {
+pad_d = param->padding[0] * 2;
+pad_h = param->padding[0] * 2;
+pad_w = param->padding[0] * 2;
+ } else if (param->padding.size() == 3) {
+// (front, top, left)
+pad_d = param->padding[0] * 2;
+pad_h = param->padding[1] * 2;
+pad_w = param->padding[2] * 2;
+ } else if (param->padding.size() == 6) {
+// (front, top, left, back, bottom, right)
+pad_d = param->padding[0] + param->padding[3];
+pad_h = param->padding[1] + param->padding[4];
+pad_w = param->padding[2] + param->padding[5];
+ } else {
+return false;
+ }
+
+ std::vector oshape;
+ for (const auto& e : dshape) {
+oshape.push_back(e);
+ }
+
+ std::vector idxes = {didx, hidx, widx};
+ for (int i = 0; i < 3; i++) {
+int ii = idxes[i];
+if (dshape[ii].as()) {
+ oshape[ii] = dshape[ii];
+} else {
+ if (param->ceil_mode) {
+oshape[ii] = ((dshape[ii] + pad_d - param->pool_size[i] +
+ param->strides[i] - 1) / param->strides[i]) + 1;
+ } else {
+oshape[ii] = ((dshape[ii] + pad_d - param->pool_size[i]) /
param->strides[i]) + 1;
+ }
+}
+ }
+
+ // assign output type
+ reporter->Assign(types[1], TensorTypeNode::make(oshape, data->dtype));
+ return true;
+}
+
+// MaxPool3D
+Expr MakeMaxPool3D(Expr data,
Review comment:
I think you can unify MakeMaxPool2D and MakeMaxPool3D with a template param
for Attrs and additional argument for op name ("nn.max_pool2d" or
"nn.max_pool3d"). Also for avg pool. Can you try 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:
us...@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r356735856
##
File path: src/relay/op/nn/pooling.cc
##
@@ -720,5 +720,238 @@ RELAY_REGISTER_OP("nn.avg_pool2d_grad")
.set_attr("FTVMCompute", Pool2DGradCompute);
+// relay.nn.max_pool3d & relay.nn.avg_pool3d
+TVM_REGISTER_NODE_TYPE(MaxPool3DAttrs);
+TVM_REGISTER_NODE_TYPE(AvgPool3DAttrs);
+
+template
+bool Pool3DRel(const Array& types,
+ int num_inputs,
+ const Attrs& attrs,
+ const TypeReporter& reporter) {
+ CHECK_EQ(types.size(), 2);
+ const auto* data = types[0].as();
+
+ if (data == nullptr) return false;
+
+ const auto dshape = data->shape;
+ CHECK_GE(dshape.size(), 3U)
+ << "Pool3D only support input >= 3-D: input must have depth, height and
width";
+ const auto param = attrs.as();
+ CHECK(param != nullptr);
+
+ Layout layout(param->layout);
+ CHECK(layout.Contains(LayoutAxis::Get('D')) &&
layout.Contains(LayoutAxis::Get('H')) &&
+layout.Contains(LayoutAxis::Get('W')) &&
!layout.Contains(LayoutAxis::Get('d')) &&
+!layout.Contains(LayoutAxis::Get('h')) &&
!layout.Contains(LayoutAxis::Get('w')))
+<< "Invalid layout " << layout
+<< ". Pool3D layout must have D, H and W, which cannot be split";
+
+ const auto didx = layout.IndexOf(LayoutAxis::Get('D'));
+ const auto hidx = layout.IndexOf(LayoutAxis::Get('H'));
+ const auto widx = layout.IndexOf(LayoutAxis::Get('W'));
+
+ IndexExpr pad_d, pad_h, pad_w;
+ if (param->padding.size() == 1) {
+pad_d = param->padding[0] * 2;
+pad_h = param->padding[0] * 2;
+pad_w = param->padding[0] * 2;
+ } else if (param->padding.size() == 3) {
+// (front, top, left)
+pad_d = param->padding[0] * 2;
+pad_h = param->padding[1] * 2;
+pad_w = param->padding[2] * 2;
+ } else if (param->padding.size() == 6) {
+// (front, top, left, back, bottom, right)
+pad_d = param->padding[0] + param->padding[3];
+pad_h = param->padding[1] + param->padding[4];
+pad_w = param->padding[2] + param->padding[5];
+ } else {
+return false;
+ }
+
+ std::vector oshape;
+ for (const auto& e : dshape) {
+oshape.push_back(e);
+ }
+
+ std::vector idxes = {didx, hidx, widx};
+ for (int i = 0; i < 3; i++) {
+int ii = idxes[i];
+if (dshape[ii].as()) {
+ oshape[ii] = dshape[ii];
+} else {
+ if (param->ceil_mode) {
+oshape[ii] = ((dshape[ii] + pad_d - param->pool_size[i] +
+ param->strides[i] - 1) / param->strides[i]) + 1;
+ } else {
+oshape[ii] = ((dshape[ii] + pad_d - param->pool_size[i]) /
param->strides[i]) + 1;
+ }
+}
+ }
+
+ // assign output type
+ reporter->Assign(types[1], TensorTypeNode::make(oshape, data->dtype));
+ return true;
+}
+
+// MaxPool3D
+Expr MakeMaxPool3D(Expr data,
Review comment:
I think you can unify MakeMaxPool2D and MakeMaxPool3D with a template param
for Attrs and additional argument for op name ("nn.max_pool2d" or
"nn.max_pool3d").
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:
us...@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r355132397
##
File path: src/relay/op/nn/pooling.cc
##
@@ -720,5 +720,257 @@ RELAY_REGISTER_OP("nn.avg_pool2d_grad")
.set_attr("FTVMCompute", Pool2DGradCompute);
+// relay.nn.max_pool3d & relay.nn.avg_pool3d
+TVM_REGISTER_NODE_TYPE(MaxPool3DAttrs);
+TVM_REGISTER_NODE_TYPE(AvgPool3DAttrs);
+
+template
+Array > Pool3DInferCorrectLayout(
Review comment:
I think this function can be unified for 2d and 3d.
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:
us...@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r355131637
##
File path: topi/src/topi.cc
##
@@ -595,7 +602,7 @@ TVM_REGISTER_GLOBAL("topi.generic.schedule_injective")
TVM_REGISTER_GLOBAL("topi.generic.schedule_injective_from_existing")
.set_body([](TVMArgs args, TVMRetValue *rv) {
*rv = topi::generic::schedule_injective_from_existing(args[0], args[1]);
- });
+ });
Review comment:
Remove this diff and other three below
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:
us...@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
URL: https://github.com/apache/incubator-tvm/pull/4478#discussion_r355131637
##
File path: topi/src/topi.cc
##
@@ -595,7 +602,7 @@ TVM_REGISTER_GLOBAL("topi.generic.schedule_injective")
TVM_REGISTER_GLOBAL("topi.generic.schedule_injective_from_existing")
.set_body([](TVMArgs args, TVMRetValue *rv) {
*rv = topi::generic::schedule_injective_from_existing(args[0], args[1]);
- });
+ });
Review comment:
Remove this diff
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:
us...@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #4478: [TOPI] implement pool3d op
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{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& kernel_size,
+ const Array& stride_size,
+ const Array& padding_size,
+ PoolType pool_type,
+ bool ceil_mode,
+ const std::vector& 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 daxis;
+ std::vector kernel(k_size);
+ std::vector stride(k_size);
+ std::vector pad_head(k_size);
+ std::vector pad_tail(k_size);
+ Array pad_before(std::vector(x_size, 0));
+ Array pad_after(std::vector(x_size, 0));
+ Array 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& output) {
+ Array 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& output) {
+ Array 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& output) {
+ Array 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 start(k_size);
+std::vector 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
