add optimizer_extrace_pass (PaddlePaddle#6)

paddle/fluid/framework/ipu/ipu_backend.h

@@ -38,17 +38,34 @@ limitations under the License. */
 namespace paddle {
 namespace framework {
 
+struct Optimizer {
+  std::string type;
+  // as far as we know, attr is usually float
+  std::map<std::string, float> attrs;
+};
+
 class IpuBackend {
  public:
   explicit IpuBackend();
 
-  void Compile(ir::Graph *graph,
-               const std::vector<std::string> &feed_list,
+  void Compile(ir::Graph *graph, const std::vector<std::string> &feed_list,
                const std::vector<std::string> &fetch_list);
 
   void Run(const std::vector<const Tensor *> &inputs,
            std::vector<Tensor *> &outputs);
 
+  std::string GetOptimizerType() { return optimizer_.type; }
+
+  void SetOptimizerType(const std::string &type) { optimizer_.type = type; }
+
+  const std::map<std::string, float> &GetOptimizerAttr() {
+    return optimizer_.attrs;
+  }
+
+  void SetOptimizerAttr(const std::string &attr, float value) {
+    optimizer_.attrs[attr] = value;
+  }
+
   static std::shared_ptr<IpuBackend> GetInstance() {
     if (NULL == instance_) {
       instance_.reset(new IpuBackend());
@@ -60,6 +77,7 @@ class IpuBackend {
   // std::map<std::string, popart::TensorId> inputs_;
   // std::map<std::string, popart::TensorId> outputs_;
   // std::map<std::string, popart::TensorId> tensors_;
+  Optimizer optimizer_;
 
   std::vector<popart::TensorId> inputs_;
   std::vector<popart::TensorId> outputs_;

paddle/fluid/framework/ir/CMakeLists.txt

@@ -126,6 +126,7 @@ endif()
 if(WITH_IPU)
     cc_library(ipu_pass_base SRCS ipu/ipu_pass_base.cc DEPS pass)
     pass_library(forward_graph_extract_pass base DIR ipu DEPS ipu_pass_base)
+    pass_library(optimizer_extract_pass base DIR ipu DEPS ipu_pass_base)
     pass_library(ipu_graph_builder_pass base DIR ipu DEPS ipu_pass_base)
     pass_library(ipu_runtime_replacer_pass base DIR ipu DEPS ipu_pass_base)
 endif()

paddle/fluid/framework/ir/ipu/optimizer_extract_pass.cc

@@ -0,0 +1,61 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/framework/ir/ipu/optimizer_extract_pass.h"
+
+#include "paddle/fluid/framework/ipu/ipu_backend.h"
+#include "paddle/fluid/framework/op_proto_maker.h"
+
+namespace paddle {
+namespace framework {
+namespace ir {
+
+class Graph;
+
+void IpuOptimizerExtractPass::ApplyImpl(ir::Graph* graph) const {
+  auto ipu_backend = paddle::framework::IpuBackend::GetInstance();
+
+  for (auto* node : graph->Nodes()) {
+    if (node->IsOp() && node->Op()) {
+      int op_role = BOOST_GET_CONST(
+          int, node->Op()->GetAttr(
+                   framework::OpProtoAndCheckerMaker::OpRoleAttrName()));
+
+      // graph usually have multiple optimizer node for different parameter,
+      // and these node have the same type and attr value usually
+      if ((op_role == static_cast<int>(framework::OpRole::kOptimize))) {
+        ipu_backend->SetOptimizerType(node->Op()->Type());
+
+        for (const std::string& attr_name : node->Op()->AttrNames()) {
+          auto attr_type = node->Op()->GetAttrType(attr_name);
+          // with adam, attr are float
+          if (attr_type == proto::AttrType::FLOAT) {
+            auto attr_value =
+                BOOST_GET_CONST(float, node->Op()->GetAttr(attr_name));
+            ipu_backend->SetOptimizerAttr(attr_name, attr_value);
+          } else {
+            VLOG(10) << "Skip " << attr_type;
+          }
+        }
+      }
+    }
+  }
+}
+
+}  // namespace ir
+}  // namespace framework
+}  // namespace paddle
+
+REGISTER_PASS(optimizer_extract_pass,
+              paddle::framework::ir::IpuOptimizerExtractPass);

paddle/fluid/framework/ir/ipu/optimizer_extract_pass.h

@@ -0,0 +1,29 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/framework/ir/graph.h"
+#include "paddle/fluid/framework/ir/ipu/ipu_pass_base.h"
+
+namespace paddle {
+namespace framework {
+namespace ir {
+
+class IpuOptimizerExtractPass : public IPUPassBase {
+ protected:
+  void ApplyImpl(ir::Graph* graph) const override;
+};
+
+}  // namespace ir
+}  // namespace framework
+}  // namespace paddle