cpu: pooling: fix failde examples

Author: ankalinin

src/cpu/cpu_pooling_pd.hpp

@@ -54,6 +54,15 @@ struct cpu_pooling_fwd_pd_t: public pooling_fwd_pd_t {
     cpu_memory_pd_t ws_pd_;
 
     virtual status_t init() = 0;
+
+    virtual status_t set_default_params() {
+        using namespace memory_format;
+        if (src_pd_.desc()->format == any)
+            CHECK(src_pd_.set_format(nchw));
+        if (dst_pd_.desc()->format == any)
+            CHECK(dst_pd_.set_format(src_pd_.desc()->format));
+        return status::success;
+    }
 };
 
 struct cpu_pooling_bwd_pd_t: public pooling_bwd_pd_t {
@@ -80,6 +89,15 @@ struct cpu_pooling_bwd_pd_t: public pooling_bwd_pd_t {
     cpu_memory_pd_t ws_pd_;
 
     virtual status_t init() = 0;
+
+    virtual status_t set_default_params() {
+        using namespace memory_format;
+        if (diff_dst_pd_.desc()->format == any)
+            CHECK(diff_dst_pd_.set_format(nchw));
+        if (diff_src_pd_.desc()->format == any)
+            CHECK(diff_src_pd_.set_format(diff_dst_pd_.desc()->format));
+        return status::success;
+    }
 };
 
 }

src/cpu/jit_avx2_pooling.hpp

@@ -44,6 +44,7 @@ struct jit_avx2_pooling_fwd_t: public cpu_primitive_t {
             using namespace utils;
             assert(engine()->kind() == engine_kind::cpu);
             bool ok = true
+                && set_default_params() == status::success
                 && one_of(desc()->prop_kind, forward_training,
                         forward_inference)
                 && one_of(desc()->alg_kind, pooling_max, pooling_avg)
@@ -98,6 +99,7 @@ struct jit_avx2_pooling_bwd_t: public cpu_primitive_t {
             using namespace utils;
             assert(engine()->kind() == engine_kind::cpu);
             bool ok = true
+                && set_default_params() == status::success
                 && one_of(desc()->prop_kind, backward, backward_data)
                 && one_of(desc()->alg_kind, pooling_max, pooling_avg)
                 && everyone_is(data_type::f32, diff_src_pd()->desc()->data_type,

src/cpu/ref_pooling.hpp

@@ -43,6 +43,7 @@ struct ref_pooling_fwd_t: public cpu_primitive_t {
             using namespace alg_kind;
             assert(engine()->kind() == engine_kind::cpu);
             bool ok = true
+                && set_default_params() == status::success
                 && utils::one_of(desc()->prop_kind, forward_training,
                         forward_inference)
                 && utils::one_of(desc()->alg_kind, pooling_max, pooling_avg)
@@ -90,6 +91,7 @@ struct ref_pooling_bwd_t: public cpu_primitive_t {
             using namespace alg_kind;
             assert(engine()->kind() == engine_kind::cpu);
             bool ok = true
+                && set_default_params() == status::success
                 && utils::one_of(desc()->prop_kind, backward_data)
                 && utils::one_of(desc()->alg_kind, pooling_max, pooling_avg)
                 && utils::everyone_is(data_type, diff_dst_pd()->desc()->data_type,