diff --git a/src/gopt/impl/padding_channel.cpp b/src/gopt/impl/padding_channel.cpp index 8261ff734abb962c8c208b265547aa60274637b1..b47fd9bbe858c4715ecf2fe11b7bc8eeb3204a6e 100644 --- a/src/gopt/impl/padding_channel.cpp +++ b/src/gopt/impl/padding_channel.cpp @@ -114,7 +114,14 @@ void PaddingChannelPass::apply(OptState& opt) const { auto rewriter = opt.graph().make_rewriter(); auto on_opr = [this, &opt, &rewriter](OperatorNodeBase* opr) { auto it = m_opr_replace_funcs.find(opr->dyn_typeinfo()); - if (it != m_opr_replace_funcs.end()) { + auto is_skip = false; + //! if the input of the opr is dynamic shape, skip it + for (size_t id = 0; id < opr->input().size(); id++) { + if (0 == opr->input(id)->shape().ndim) { + is_skip = true; + } + } + if (it != m_opr_replace_funcs.end() && !is_skip) { VarNodeArray new_inp; new_inp.reserve(opr->input().size()); for (auto&& inp : opr->input()) { diff --git a/src/gopt/test/padding_channel.cpp b/src/gopt/test/padding_channel.cpp index 58f734e434383bd46597368d981af5a1819a71f0..e7dd1493fc19b172a2f592a6aee7a6f275da58c8 100644 --- a/src/gopt/test/padding_channel.cpp +++ b/src/gopt/test/padding_channel.cpp @@ -443,4 +443,59 @@ TEST(TestGoptInference, ChannelPaddingMoreOp) { MGB_ASSERT_TENSOR_NEAR(host_y, host_y_opt, 1e-3); } +TEST(TestGoptInference, DynamicShape) { + HostTensorGenerator<> gen; + HostTensorGenerator gen_int; + auto cn = CompNode::load("cpu0"); + auto graph = ComputingGraph::make(); + graph->options().graph_opt_level = 0; + auto mkcvar = [&](const char* name, const TensorShape& shp) { + return opr::SharedDeviceTensor::make(*graph, *gen(shp, cn)).rename(name); + }; + + auto host_x = gen({1, 4, 8, 8}, cn); + auto host_dy = gen_int({1}, cn); + auto x = opr::Host2DeviceCopy::make(*graph, host_x); + auto dy = opr::Host2DeviceCopy::make(*graph, host_dy); + int32_t* start = host_dy->ptr(); + start[0] = 0; + dy = opr::MarkDynamicVar::make(dy); + using AIdx = opr::indexing::AxisIndexer; + auto sub = opr::Subtensor::make( + x, {AIdx::make_interval(1, dy, x.make_scalar(4), x.make_scalar(1))}, + OperatorNodeConfig("sub")); + + //! Hybrid nchw44 mode + opr::ConvBias::Param param_conv; + param_conv.pad_h = param_conv.pad_w = 1; + auto w1 = mkcvar("w1", {3, 4, 3, 3}), b1 = mkcvar("w1", {1, 3, 1, 1}), + conv1 = opr::ConvBias::make( + sub, w1, b1, param_conv, {}, OperatorNodeConfig("conv1")); + + auto w2 = mkcvar("w2", {4, 3, 1, 1}), + y = opr::Convolution::make(conv1, w2, {}, {}, OperatorNodeConfig("conv2")); + + SymbolVar y_pad; + unpack_vector( + gopt::GraphOptimizer{} + .add_pass(gopt::PaddingChannelPass::make( + cg::GraphCommonOptimizeOptions::LayoutTransform::NCHW44, + true)) + .apply({{y}}) + .endpoint_vars(), + y_pad); + auto conv1_opt = find_opr(y_pad, "conv1"); + auto conv2_opt = find_opr(y_pad, "conv2"); + //! do not padding input tensor + ASSERT_EQ(conv1_opt->input(0)->shape().ndim, 0); + //! output tensor padding input tensor + ASSERT_EQ(conv2_opt->input(0)->shape().ndim, 0); + + HostTensorND host_y_opt, host_y; + auto func = graph->compile( + {make_callback_copy(y, host_y), make_callback_copy(y_pad, host_y_opt)}); + func->execute(); + MGB_ASSERT_TENSOR_NEAR(host_y, host_y_opt, 1e-2); +} + // vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}