cpu: concat: include padded dims when computing the physical tensor size

this closes #606
oneapi-src · Nov 29, 2019 · a265c7d · a265c7d
1 parent 0d85005
commit a265c7d
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 2 deletions.
diff --git a/src/cpu/simple_concat.cpp b/src/cpu/simple_concat.cpp
@@ -62,7 +62,7 @@ status_t simple_concat_t<data_type>::execute(const exec_ctx_t &ctx) const {
     dims_t phys_dims;
     for (size_t i = 0; i < sizeof(phys_dims) / sizeof(phys_dims[0]); i++)
         phys_dims[i] = (i < (size_t)perm[concat_dim])
-                ? o_d.dims()[iperm[i]] / pd()->blocks_[iperm[i]]
+                ? o_d.padded_dims()[iperm[i]] / pd()->blocks_[iperm[i]]
                 : 1;
 
     if (perm[concat_dim] == 0) {

diff --git a/src/cpu/simple_concat.hpp b/src/cpu/simple_concat.hpp
@@ -109,7 +109,7 @@ struct simple_concat_t : public primitive_impl_t {
 
             dim_t nelems = 1;
             for (int i = perm_[concat_dim()]; i < ndims; i++)
-                nelems *= data_d.dims()[iperm_[i]] / blocks_[iperm_[i]];
+                nelems *= data_d.padded_dims()[iperm_[i]] / blocks_[iperm_[i]];
             for (int i = 0; i < ndims; i++)
                 nelems *= blocks_[i];
 

diff --git a/tests/benchdnn/inputs/concat/test_concat_all b/tests/benchdnn/inputs/concat/test_concat_all
@@ -25,5 +25,15 @@
 6x8x3x4:6x1x3x4:6x0x3x4:6x3x3x4
 6x0x3x4:6x3x3x4:6x5x3x4:6x5x3x4
 
+# sizes are not multiple of blocks + (non-blocking) axis
+--sdt=f32,s32,s8
+--ddt=f32,s32,s8
+--dtag=undef,nhwc,nChw16c
+--axis=2
+--stag=nChw16c:nChw16c
+6x5x3x4:6x5x3x4
+6x25x3x4:6x25x3x4
+6x23x0x4:6x23x3x4
+
 # bf16
 --batch=test_concat_bfloat16