From 4ea19f51ee84c068b9270765704b34205325b1aa Mon Sep 17 00:00:00 2001
From: yuchaojie <yuchaojie1@huawei.com>
Date: Thu, 17 Dec 2020 09:59:00 +0800
Subject: [PATCH] add MaxPool2D/1D unify_mindir

---
 .../mindir/maxpool_with_argmax_unify_mindir.cc |  4 ----
 .../ccsrc/backend/session/ascend_session.cc    |  1 +
 mindspore/nn/layer/pooling.py                  | 18 ++----------------
 3 files changed, 3 insertions(+), 20 deletions(-)

diff --git a/mindspore/ccsrc/backend/optimizer/ascend/mindir/maxpool_with_argmax_unify_mindir.cc b/mindspore/ccsrc/backend/optimizer/ascend/mindir/maxpool_with_argmax_unify_mindir.cc
index 5f2cd4722a..7d7d8db6c5 100644
--- a/mindspore/ccsrc/backend/optimizer/ascend/mindir/maxpool_with_argmax_unify_mindir.cc
+++ b/mindspore/ccsrc/backend/optimizer/ascend/mindir/maxpool_with_argmax_unify_mindir.cc
@@ -73,8 +73,6 @@ const AnfNodePtr MaxPoolWithArgmaxUnifyMindIR::Process(const FuncGraphPtr &graph
   auto shapes = {output_shape, argmax_shape};
   AnfAlgo::SetOutputInferTypeAndShape(types, shapes, maxpool_with_argmax.get());
 
-  auto manager = graph->manager();
-  MS_EXCEPTION_IF_NULL(manager);
   return maxpool_with_argmax;
 }
 
@@ -107,8 +105,6 @@ const AnfNodePtr MaxPoolGradWithArgmaxUnifyMindIR::Process(const FuncGraphPtr &g
   argmax_shape[2] = ksize[1] * ksize[2];
   AnfAlgo::SetOutputInferTypeAndShape({argmax_dtype}, {argmax_shape}, tuple_getitem0_anf.get());
 
-  auto manager = graph->manager();
-  MS_EXCEPTION_IF_NULL(manager);
   return maxpool_grad_with_argmax;
 }
 }  // namespace opt
diff --git a/mindspore/ccsrc/backend/session/ascend_session.cc b/mindspore/ccsrc/backend/session/ascend_session.cc
index 3ac12624a0..0fa8926da7 100644
--- a/mindspore/ccsrc/backend/session/ascend_session.cc
+++ b/mindspore/ccsrc/backend/session/ascend_session.cc
@@ -439,6 +439,7 @@ void AscendSession::UnifyMindIR(const KernelGraphPtr &graph) {
   }
   auto optimizer = std::make_shared<opt::GraphOptimizer>();
   auto unify_mindir_pm = std::make_shared<opt::PassManager>("unify_mindir_pm");
+  unify_mindir_pm->AddPass(std::make_shared<opt::MaxPool2MaxPoolWithArgmax>());
   unify_mindir_pm->AddPass(std::make_shared<opt::MaxPoolWithArgmaxUnifyMindIR>());
   unify_mindir_pm->AddPass(std::make_shared<opt::MaxPoolGradWithArgmaxUnifyMindIR>());
   unify_mindir_pm->AddPass(std::make_shared<opt::Conv2DUnifyMindIR>());
diff --git a/mindspore/nn/layer/pooling.py b/mindspore/nn/layer/pooling.py
index d5ebb785bf..22841eca64 100644
--- a/mindspore/nn/layer/pooling.py
+++ b/mindspore/nn/layer/pooling.py
@@ -124,16 +124,9 @@ class MaxPool2d(_PoolNd):
                                   strides=self.stride,
                                   padding=self.pad_mode,
                                   data_format=self.format)
-        self.max_pool_with_arg_max = P.MaxPoolWithArgmax(ksize=self.kernel_size,
-                                                         strides=self.stride,
-                                                         padding=self.pad_mode)
-        self.is_tbe = context.get_context("device_target") == "Ascend"
 
     def construct(self, x):
-        if self.is_tbe and self.training:
-            out = self.max_pool_with_arg_max(x)[0]
-        else:
-            out = self.max_pool(x)
+        out = self.max_pool(x)
         return out
 
 
@@ -198,22 +191,15 @@ class MaxPool1d(_PoolNd):
         self.max_pool = P.MaxPool(ksize=self.kernel_size,
                                   strides=self.stride,
                                   padding=self.pad_mode)
-        self.max_pool_with_arg_max = P.MaxPoolWithArgmax(ksize=self.kernel_size,
-                                                         strides=self.stride,
-                                                         padding=self.pad_mode)
         self.shape = F.shape
         self.reduce_mean = P.ReduceMean(keep_dims=True)
         self.expand = P.ExpandDims()
         self.squeeze = P.Squeeze(2)
-        self.is_tbe = context.get_context("device_target") == "Ascend"
 
     def construct(self, x):
         _shape_check(self.shape(x))
         x = self.expand(x, 2)
-        if self.is_tbe and self.training:
-            output = self.max_pool_with_arg_max(x)[0]
-        else:
-            output = self.max_pool(x)
+        output = self.max_pool(x)
         output = self.squeeze(output)
         return output