- Added test for max pooling to verify mkldnn maxpool implementation

- added workaround to attach the maxpool workspace for bprop delta propogation

src/ngraph/runtime/cpu/cpu_emitter.cpp

@@ -2536,7 +2536,6 @@ void runtime::cpu::CPU_Emitter::EMITTER_DECL(EmitMaxPoolBackprop)
 
         writer << "{\n";
         writer.indent++;
-
         writer << "engine cpu_engine = engine(engine::cpu, 0);\n";
         writer << "memory::desc input_data_desc = memory::desc({" << join(delta_shape) << "}, "
                << et << ", memory::format::nchw);\n";
@@ -2547,26 +2546,38 @@ void runtime::cpu::CPU_Emitter::EMITTER_DECL(EmitMaxPoolBackprop)
         writer << "memory result = memory({result_desc, cpu_engine}, " << out[0].get_name()
                << ");\n";
 
+        //----------------------------------------------------------------------------------------------
         // create a forward primitive_desc, use this to query the workspace
-        // TODO: we need to develop global context to keep the mapping of fprop annd bprop corrosponding
-        // mkldnn kernels and use it to query the workspace requirement during bprop
-        // Note: 
-        //      input_data_desc of MaxpoolBackprop : will be same as maxpool(fprop) result_desc
-        //      result_desc of MaxpoolBackprop : will be same as maxpool(fprop) input_data_desc
+        // FIXME: (pruthvi) this is a workaround, till we maintain a global context to refer to the corrosponding
+        //        MKLDNN kernel. 
+        writer << "memory::desc max_pool_input_desc = memory::desc({" << join(args[0].get_shape()) << "}, " << et
+               << ", memory::format::nchw);\n";
+        writer << "memory::desc max_pool_result_desc = memory::desc({" << join(args[1].get_shape()) << "}, " << et
+               << ", memory::format::nchw);\n";
+        writer << "memory maxpool_input_data = memory({max_pool_input_desc, cpu_engine}, " << args[0].get_name()
+               << ");\n";
+        writer << "memory maxpool_result = memory({max_pool_result_desc, cpu_engine}, " << out[0].get_name()
+               << ");\n";
         writer << "pooling_forward::primitive_desc pool_fwd_pd = pooling_forward::primitive_desc("
                << "{prop_kind::forward, algorithm::pooling_max, "
-               << "result_desc, input_data_desc, {" << join(max_pool_fprop_op->get_window_movement_strides())
+               << "max_pool_input_desc, max_pool_result_desc, {" << join(max_pool_fprop_op->get_window_movement_strides())
                << "}, {" << join(max_pool_fprop_op->get_window_shape()) << "}, "
                << "{" << join(max_pool_fprop_op->get_padding_below()) << "}, "
                << "{" << join(max_pool_fprop_op->get_padding_above()) << "}, "
                << "padding_kind::zero}, cpu_engine);\n";
 
-
-        // query the workspace from the forward primitive desc
-        writer << "memory max_pool_workspace_memory = "
+        // query the workspace from the forward primitive desc and allocates memory
+        writer << "auto max_pool_workspace_memory = "
                   "memory(pool_fwd_pd.workspace_primitive_desc());\n";
+        //run fprop with this workspace attached
+        writer << "pooling_forward max_pooling_fwd = pooling_forward("
+               << "pool_fwd_pd, maxpool_input_data, maxpool_result, max_pool_workspace_memory);\n";
 
-        writer << "auto avg_pooling = pooling_backward(pooling_backward::primitive_desc("
+        writer << "stream s_fprop = stream(stream::kind::eager);\n"
+               << "s_fprop.submit({max_pooling_fwd}).wait();\n";
+
+        //---------------------------------------------------------------------------------------------
+        writer << "auto max_pooling_bwd = pooling_backward(pooling_backward::primitive_desc("
                << "pooling_backward::desc(algorithm::pooling_max, "
                << "result_desc, input_data_desc, {" << join(mpb->get_window_movement_strides())
                << "}, {" << join(mpb->get_window_shape()) << "}, "
@@ -2574,8 +2585,9 @@ void runtime::cpu::CPU_Emitter::EMITTER_DECL(EmitMaxPoolBackprop)
                << "{" << join(mpb->get_padding_above()) << "}, "
                << "padding_kind::zero), cpu_engine, pool_fwd_pd), "
                << "input_data, max_pool_workspace_memory, result);\n";
-        writer << "auto s = stream(stream::kind::eager);\n"
-               << "s.submit({avg_pooling}).wait();\n";
+        writer << "auto s_bwd = stream(stream::kind::eager);\n"
+               << "s_bwd.submit({max_pooling_bwd}).wait();\n";
+
         writer.indent--;
         writer << "}\n";
     }

test/autodiff.in.cpp

@@ -1405,3 +1405,94 @@ TEST(${BACKEND_NAME}, backwards_reverse_3d_02)
     };
     EXPECT_TRUE(autodiff_numeric_compare<float>(manager, backend, make_graph, {x}, .01f, .01f));
 }
+
+TEST(${BACKEND_NAME}, backwards_maxpool_n4c1h4w4_kh2kw2_sh1sw1)
+{
+    auto manager = runtime::Manager::get("${BACKEND_NAME}");
+    auto backend = manager->allocate_backend();
+    pass::Manager pass_manager;
+    Shape shape_a{1, 4, 4, 4}; //in CHWN
+    Shape maxpool_shape{4, 1, 3, 3};
+
+    auto A = make_shared<op::Parameter>(element::f32, shape_a);
+    auto reshape = make_shared<op::Reshape>(
+        A, AxisVector{0, 3, 1, 2}, Shape{4, 1, 4, 4}); //convert CHWN to NCHW
+    Shape window_shape{2, 2};
+    auto window_movement_strides = Strides{1, 1};
+    auto maxpool = make_shared<op::MaxPool>(reshape, window_shape, window_movement_strides);
+    auto f = make_shared<Function>(maxpool, op::Parameters{A});
+    shared_ptr<runtime::TensorView> ep =
+        backend->make_primary_tensor_view(element::f32, maxpool_shape);
+    vector<float> dataEp(shape_size(maxpool_shape), 4);
+
+    shared_ptr<runtime::TensorView> input =
+        backend->make_primary_tensor_view(element::f32, shape_a);
+    shared_ptr<runtime::TensorView> output =
+        backend->make_primary_tensor_view(element::f32, shape_a);
+
+    vector<float> dataInput{11, 65, 44, 28, 31, 33, 21, 66, 40, 49, 69, 57, 47, 30, 24, 27,
+                          13, 56, 46, 60, 61, 41, 25, 42, 48, 53, 51, 43, 59, 58, 29, 71,
+                          17, 22, 72, 18, 39, 35, 15, 38, 64, 52, 73, 67, 62, 50, 10, 68,
+                          45, 63, 16, 14, 55, 54, 37, 20, 36, 12, 70, 34, 19, 26, 32, 23};
+
+    vector<float> expected{//delta
+                         0, 4, 0, 0, 0, 0, 0, 8, 0, 0, 8, 0, 0, 0, 0, 0, 0, 4, 4,  4, 12, 0,
+                         0, 0, 0, 8, 0, 0, 4, 8, 0, 8, 0, 0, 8, 0, 0, 0, 0, 4, 16, 4, 16, 8,
+                         0, 0, 0, 4, 0, 4, 0, 0, 4, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0,  0};
+
+    copy_data(ep, dataEp);
+    copy_data(input, dataInput);
+
+    auto C = make_shared<op::Parameter>(element::f32, maxpool_shape);
+    auto df = autodiff::backprop_function(f);
+    pass_manager.run_passes(df);
+    auto external = manager->compile(df);
+    auto cf = backend->make_call_frame(external);
+    cf->tensor_call({input, ep}, {output});
+
+    ASSERT_TRUE(read_vector<float>(output) == expected);
+}
+
+TEST(${BACKEND_NAME}, backwards_maxpool_n2c1h5w5_kh3kw3_sh2sw2)
+{
+    auto manager = runtime::Manager::get("${BACKEND_NAME}");
+    auto backend = manager->allocate_backend();
+
+    Shape shape_a{1, 5, 5, 2}; //in CHWN
+    Shape maxpool_shape{1, 2, 2, 2};
+
+    auto A = make_shared<op::Parameter>(element::f32, shape_a);
+    auto reshape = make_shared<op::Reshape>(
+        A, AxisVector{0, 3, 1, 2}, Shape{1, 2, 5, 5}); //convert CHWN to NCHW
+    Shape window_shape{3, 3};
+    auto window_movement_strides = Strides{2, 2};
+    auto maxpool = make_shared<op::MaxPool>(reshape, window_shape, window_movement_strides);
+    auto f = make_shared<Function>(maxpool, op::Parameters{A});
+
+    shared_ptr<runtime::TensorView> ep =
+        backend->make_primary_tensor_view(element::f32, maxpool_shape);
+    vector<float> dataEp(shape_size(maxpool_shape), 4);
+
+    shared_ptr<runtime::TensorView> input =
+        backend->make_primary_tensor_view(element::f32, shape_a);
+    shared_ptr<runtime::TensorView> output =
+        backend->make_primary_tensor_view(element::f32, shape_a);
+
+    vector<float> dataInput{58, 15, 51, 35, 18, 47, 31, 32, 52, 21, 36, 38, 57, 54, 25, 45, 23,
+                          30, 16, 27, 48, 20, 41, 37, 43, 39, 22, 28, 33, 29, 12, 17, 44, 42,
+                          19, 40, 10, 46, 34, 53, 26, 55, 50, 13, 24, 14, 49, 56, 59, 11};
+
+    vector<float> expected{//delta
+                         4, 0, 0, 0, 0, 4, 0, 0, 4, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 4, 0, 0, 0, 0, 4, 4, 0};
+
+    copy_data(ep, dataEp);
+    copy_data(input, dataInput);
+
+    auto C = make_shared<op::Parameter>(element::f32, maxpool_shape);
+    auto df = autodiff::backprop_function(f);
+    auto external = manager->compile(df);
+    auto cf = backend->make_call_frame(external);
+    cf->tensor_call({input, ep}, {output});
+    ASSERT_TRUE(read_vector<float>(output) == expected);
+}
\ No newline at end of file