From eb826422c404e4628fb2714ab22b9dd59681b87f Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Tue, 15 Dec 2020 18:47:23 +0800
Subject: [PATCH] fix(dnn): forbid pooling window size smaller than padding

GitOrigin-RevId: 9ad61c409d741ded70c88473374e1de426b27008
---
 dnn/src/common/pooling.cpp   |  4 ++++
 src/opr/test/dnn/pooling.cpp | 11 +++++++++++
 2 files changed, 15 insertions(+)

diff --git a/dnn/src/common/pooling.cpp b/dnn/src/common/pooling.cpp
index ef74bfd6..634aed4d 100644
--- a/dnn/src/common/pooling.cpp
+++ b/dnn/src/common/pooling.cpp
@@ -92,6 +92,10 @@ void PoolingBase::deduce_layout_fwd(const TensorLayout& src,
     size_t sw = this->param().stride_w;
     size_t ph = this->param().pad_h;
     size_t pw = this->param().pad_w;
+    megdnn_assert(ph < fh && pw < fw,
+                  "pooling padding size (%zu %zu) should not be bigger than "
+                  "window size (%zu %zu)",
+                  pw, ph, fw, fh);
     infer_conv_shape2d(ih, iw, fh, fw, sh, sw, ph, pw, oh, ow);
     if (param().format == Param::Format::NCHW) {
         dst = TensorLayout(TensorShape({n, c, oh, ow}), src.dtype);
diff --git a/src/opr/test/dnn/pooling.cpp b/src/opr/test/dnn/pooling.cpp
index 8defe15b..16ffe1a3 100644
--- a/src/opr/test/dnn/pooling.cpp
+++ b/src/opr/test/dnn/pooling.cpp
@@ -104,6 +104,17 @@ TEST(TestOprDNN, PoolingBackward)
     }
 }
 
+TEST(TestOprDNN, PoolingForwardPadding) {
+    auto graph = ComputingGraph::make();
+    Param param(Param::Mode::MAX, 2, 2, 2, 2, 2, 2);
+    SymbolVarArray symbol_inputs;
+    HostTensorGenerator<> gen;
+    auto host_tensor = gen({2, 3, 23, 24});
+    symbol_inputs.push_back(
+            mgb::opr::Host2DeviceCopy::make(*graph, host_tensor));
+    ASSERT_THROW(opr::Pooling::make(symbol_inputs[0], param), MegDNNError);
+}
+
 } // anonymous namespace
 
 // vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}
-- 
GitLab