common: fix zero-padding for tensors of small rank

src/common/memory_zero_pad.cpp

@@ -14,7 +14,7 @@
 * limitations under the License.
 *******************************************************************************/
 
-#include <assert.h>
+#include <cassert>
 
 #include "mkldnn_thread.hpp"
 #include "mkldnn_traits.hpp"
@@ -61,12 +61,14 @@ void typed_zero_pad_blk(
     const int c_tail_s = C_blocked ? dims[2] % blksize : 0;
     assert(a_tail_s || b_tail_s || c_tail_s);
 
+    const int ndims = m_d.ndims();
+    assert(1 <= ndims && ndims <= 6);
     const int A = A_blocked ? pdims[0] / blksize : dims[0];
-    const int B = B_blocked ? pdims[1] / blksize : dims[1];
-    const int C = C_blocked ? pdims[2] / blksize : dims[2];
-    const int D = m_d.ndims() > 3 ? dims[3] : 1;
-    const int E = m_d.ndims() > 4 ? dims[4] : 1;
-    const int F = m_d.ndims() > 5 ? dims[5] : 1;
+    const int B = ndims <= 1 ? 1 : B_blocked ? pdims[1] / blksize : dims[1];
+    const int C = ndims <= 2 ? 1 : C_blocked ? pdims[2] / blksize : dims[2];
+    const int D = ndims <= 3 ? 1 : dims[3];
+    const int E = ndims <= 4 ? 1 : dims[4];
+    const int F = ndims <= 5 ? 1 : dims[5];
     const int inner_blk = blk.inner_nblks == 3 ? blk.inner_blks[2] : 1;
 
     auto zeroize_tail = [&](data_t *d, const int tail_s) {