src, tests: softmax: fix incorrect outer_stride variable

Previous logic didn't account for situations when neighbor dimension
could have larger stride than it should be, like for acbd format.

Author: dzarukin

src/common/softmax_pd.hpp

@@ -81,10 +81,6 @@ struct softmax_pd_t : public primitive_desc_t {
                 dst_desc().dims + axis() + 1, ndims() - 1 - axis());
     }
 
-    dim_t outer_stride() const {
-        const memory_desc_wrapper dst_d(dst_desc());
-        return axis() > 0 ? dst_d.blocking_desc().strides[axis() - 1] : 1;
-    }
     dim_t axis_stride() const {
         const memory_desc_wrapper dst_d(dst_desc());
         return dst_d.blocking_desc().strides[axis()];

src/common/verbose.cpp

@@ -1411,7 +1411,7 @@ std::string init_info_softmax(const engine_t *e, const pd_t *pd) {
 
     ss << "src_" << md2fmt_str(src_md, pd->invariant_src_user_format_kind());
     ss << " dst_" << dst_md;
-    if (diff_dst_md) ss << " diff_dst_" << diff_dst_md;
+    if (!types::is_zero_md(diff_dst_md)) ss << " diff_dst_" << diff_dst_md;
 
     ss << "," << pd->attr() << ",";
     ss << "alg:" << pd->alg_kind() << " axis:" << pd->axis() << ",";

src/cpu/ref_softmax.cpp

@@ -54,13 +54,16 @@ status_t ref_softmax_fwd_t::execute_forward_dense(const exec_ctx_t &ctx) const {
     const memory_desc_wrapper dst_d(pd()->dst_md());
 
     const auto interim_dt = data_type::f32;
-    const dim_t ou_stride = pd()->outer_stride();
     const auto is_inplace = (src == dst);
     const auto has_padding = is_padding(dst_d);
     const auto zero_padding = has_padding && !is_inplace;
     const auto axis = pd()->axis();
     const auto axis_size = pd()->axis_size(true);
-    const auto axis_blk_size = src_d.padded_dims()[axis] - src_d.dims()[axis];
+    // Since dense implementation assumes `inner_size == 1`, and src is dense
+    // and identical to dst, outer_stride should coincide with axis_size. This
+    // allows to shuffle outer dimensions and not relying on a stride of a
+    // previous dimension.
+    const auto ou_stride = axis_size;
     const auto src_dt_size = types::data_type_size(pd()->src_md()->data_type);
     const auto dst_dt_size = types::data_type_size(pd()->dst_md()->data_type);
 
@@ -178,8 +181,9 @@ status_t ref_softmax_fwd_t::execute_forward_dense(const exec_ctx_t &ctx) const {
             io::store_float_value(dst_d.data_type(), val, dst_data, c);
         }
         if (zero_padding) {
+            const auto tail = src_d.padded_dims()[axis] - src_d.dims()[axis];
             PRAGMA_OMP_SIMD()
-            for (int i = 0; i < axis_blk_size; i++)
+            for (int i = 0; i < tail; i++)
                 io::store_float_value(
                         dst_d.data_type(), 0, dst_data, channels_ + i);
         }
@@ -315,7 +319,11 @@ status_t ref_softmax_bwd_t::execute_backward_dense(
     const memory_desc_wrapper diff_dst_d(pd()->diff_dst_md());
     const memory_desc_wrapper diff_src_d(pd()->diff_src_md());
 
-    const auto ou_stride = pd()->outer_stride();
+    // Since dense implementation assumes `inner_size == 1`, and src is dense
+    // and identical to dst, outer_stride should coincide with axis_size. This
+    // allows to shuffle outer dimensions and not relying on a stride of a
+    // previous dimension.
+    const auto ou_stride = pd()->axis_size();
 
     parallel_nd(outer_size_, [&](dim_t ou) {
         float sbr = 0;

tests/benchdnn/inputs/softmax/test_softmax_all

@@ -55,6 +55,8 @@
 --axis=0
 --batch=shapes_large
 
+--reset --stag=acbd --dtag=acbd --sdt=f32 --ddt=f32 --axis=3 1x16x384x384_n"neighbor_dim_to_axis_has_larger_stride"
+
 --batch=test_softmax_bfloat16
 
 --batch=test_softmax_float16

tests/benchdnn/inputs/softmax/test_softmax_bfloat16

@@ -44,3 +44,5 @@
 --axis=0,1
 --batch=set_0d
 --batch=shapes_2d
+
+--reset --stag=acbd --dtag=acbd --sdt=bf16 --ddt=bf16 --axis=3 1x16x384x384_n"neighbor_dim_to_axis_has_larger_stride"

tests/benchdnn/inputs/softmax/test_softmax_float16

@@ -39,3 +39,5 @@
 --axis=0,1
 --batch=set_0d
 --batch=shapes_2d
+
+--reset --stag=acbd --dtag=acbd --sdt=f16 --ddt=f16 --axis=3 1x16x384x384_n"neighbor_dim_to_axis_has_larger_stride"