From ab198b45b5b67a7d86b6cfb123fe2ab815adbc8c Mon Sep 17 00:00:00 2001
From: Ruibiao Chen <chenruibiao@baidu.com>
Date: Tue, 26 Jul 2022 12:57:27 +0800
Subject: [PATCH] Set more attrs in ReplaceScaleLossGradOp (#44576)

* Set more attrs in ReplaceScaleLossGradOp

* Fix typos

* Fix CI errors

* Add UT
---
 .../details/scale_loss_grad_op_handle.h        |  2 ++
 paddle/fluid/framework/ir/CMakeLists.txt       |  2 +-
 paddle/fluid/framework/ir/graph_helper.cc      | 18 ++++++++++++++++--
 .../fluid/tests/custom_op/CMakeLists.txt       |  2 ++
 4 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/framework/details/scale_loss_grad_op_handle.h b/paddle/fluid/framework/details/scale_loss_grad_op_handle.h
index 88c87472d03..d447d0284d8 100644
--- a/paddle/fluid/framework/details/scale_loss_grad_op_handle.h
+++ b/paddle/fluid/framework/details/scale_loss_grad_op_handle.h
@@ -44,6 +44,8 @@ struct ScaleLossGradOpHandle : public OpHandleBase {
 
   ~ScaleLossGradOpHandle() final;
 
+  proto::VarType::Type DType() const { return out_dtype_; }
+
   std::string Name() const override;
 
   platform::Place GetPlace() const { return place_; }
diff --git a/paddle/fluid/framework/ir/CMakeLists.txt b/paddle/fluid/framework/ir/CMakeLists.txt
index d31555bf724..32131703524 100755
--- a/paddle/fluid/framework/ir/CMakeLists.txt
+++ b/paddle/fluid/framework/ir/CMakeLists.txt
@@ -67,7 +67,7 @@ cc_library(
 cc_library(
   graph_helper
   SRCS graph_helper.cc
-  DEPS graph)
+  DEPS graph scale_loss_grad_op_handle)
 cc_library(
   pass
   SRCS pass.cc
diff --git a/paddle/fluid/framework/ir/graph_helper.cc b/paddle/fluid/framework/ir/graph_helper.cc
index 8afd20d6a00..d6bd2e4a80d 100644
--- a/paddle/fluid/framework/ir/graph_helper.cc
+++ b/paddle/fluid/framework/ir/graph_helper.cc
@@ -18,6 +18,7 @@ limitations under the License. */
 #include <stack>
 
 #include "paddle/fluid/framework/details/multi_devices_helper.h"
+#include "paddle/fluid/framework/details/scale_loss_grad_op_handle.h"
 #include "paddle/fluid/framework/op_proto_maker.h"
 
 DECLARE_bool(convert_all_blocks);
@@ -469,11 +470,23 @@ void RemoveControlDepInputAndOuput(OpDesc *op_desc) {
 
 static OpDesc *ReplaceScaleLossGradOp(const Node &node, OpDesc *desc) {
   desc->SetType("fill_constant");
+  desc->SetAttr("shape", std::vector<int64_t>({1}));
+  desc->SetAttr("value", 1.0f);
+
+  if (node.IsWrappedBy<details::OpHandleBase>()) {
+    details::OpHandleBase &op_hander =
+        const_cast<Node *>(&node)->Wrapper<details::OpHandleBase>();
+    desc->SetAttr(
+        "dtype",
+        dynamic_cast<details::ScaleLossGradOpHandle *>(&op_hander)->DType());
+  }
+
+  desc->SetAttr("force_cpu", false);
   desc->SetAttr(
       OpProtoAndCheckerMaker::OpRoleAttrName(),
       (static_cast<int>(OpRole::kBackward) | static_cast<int>(OpRole::kLoss)));
-  desc->SetAttr("value", 1.0f);
-  desc->SetAttr("shape", std::vector<int64_t>({1}));
+  // TODO(Ruibiao) : Set OpDeviceAttrName when needed
+
   std::vector<std::string> output_names;
   for (auto out : node.outputs) {
     output_names.emplace_back(out->Name());
@@ -503,6 +516,7 @@ static void GetGraphOpDesc(const std::vector<Node *> &nodes,
 
     // create fill_constant op
     if (n->Name() == "scale_loss_grad") {
+      VLOG(4) << "convert op node scale_loss_grad to desc fill_constant";
       ops->emplace_back();
       auto &desc = ops->back();
       ReplaceScaleLossGradOp(*n, &desc);
diff --git a/python/paddle/fluid/tests/custom_op/CMakeLists.txt b/python/paddle/fluid/tests/custom_op/CMakeLists.txt
index f21fc730fc8..2addead40fc 100644
--- a/python/paddle/fluid/tests/custom_op/CMakeLists.txt
+++ b/python/paddle/fluid/tests/custom_op/CMakeLists.txt
@@ -7,6 +7,8 @@ if(WITH_GPU OR APPLE)
 
   # Compiling shared library will cost some time, but running process is very fast.
   set_tests_properties(test_custom_relu_op_setup PROPERTIES TIMEOUT 250)
+  set_tests_properties(test_custom_relu_op_setup
+                       PROPERTIES ENVIRONMENT FLAGS_CONVERT_GRAPH_TO_PROGRAM=1)
   set_tests_properties(test_custom_relu_op_jit PROPERTIES TIMEOUT 180)
   set_tests_properties(test_custom_relu_model PROPERTIES TIMEOUT 180)
   set_tests_properties(test_context_pool PROPERTIES TIMEOUT 180)
-- 
GitLab