[PIR]Adapt no_grad api in pir mode (#67966)

0x45f · web-flow · commit 1b4b4d1c1100 · 2024-09-26T16:25:32.000+08:00
diff --git a/paddle/fluid/imperative/tracer.cc b/paddle/fluid/imperative/tracer.cc
@@ -44,15 +44,15 @@ namespace paddle {
 namespace imperative {
 thread_local std::string Tracer::python_stack_ = "";
 
-thread_local bool Tracer::has_grad_ = true;
-
 thread_local bool Tracer::use_layout_autotune_ = false;
 
 static thread_local std::shared_ptr<Tracer> g_current_tracer(nullptr);
 
 static thread_local std::shared_ptr<AmpAttrs> g_current_amp_attrs =
     std::make_shared<AmpAttrs>();
 
+static thread_local bool g_has_grad = true;
+
 TEST_API void Tracer::DisableLayoutAutoTune() { use_layout_autotune_ = false; }
 TEST_API void Tracer::EnableLayoutAutoTune() {
   use_layout_autotune_ = true;
@@ -417,7 +417,7 @@ void Tracer::TraceOp(const std::string& type,
                    outs,
                    std::move(attrs),
                    expected_place_,
-                   has_grad_,
+                   g_has_grad,
                    inplace_map);
 }
 
@@ -546,9 +546,9 @@ void Tracer::TraceOp(const std::string& type,
 TEST_API void Tracer::SetExpectedPlace(phi::Place place) {
   expected_place_ = place;
 }
-TEST_API bool Tracer::HasGrad() const { return has_grad_; }
+TEST_API bool Tracer::HasGrad() const { return g_has_grad; }
 
-TEST_API void Tracer::SetHasGrad(bool has_grad) { has_grad_ = has_grad; }
+TEST_API void Tracer::SetHasGrad(bool has_grad) { g_has_grad = has_grad; }
 
 TEST_API void Tracer::SetUsePromote(bool use_promote) {
   VLOG(4) << "set use_promote to " << use_promote;
diff --git a/paddle/fluid/imperative/tracer.h b/paddle/fluid/imperative/tracer.h
@@ -182,7 +182,6 @@ class Tracer {
   static thread_local std::string python_stack_;
   static thread_local bool enable_program_desc_tracing_;
   static thread_local bool use_layout_autotune_;
-  static thread_local bool has_grad_;
 };
 
 // To access static variable current_tracer
diff --git a/paddle/fluid/pir/dialect/op_generator/api_gen.py b/paddle/fluid/pir/dialect/op_generator/api_gen.py
@@ -147,8 +147,14 @@
     {set_null_type}
     {handle_optional_outputs}
     {out_split}
+    {set_stop_gradient}
     {return_result}"""
 
+SET_STOP_GRADIENT_TEMPLATE = """
+    if (!egr::Controller::Instance().HasGrad()) {{
+        SetStopGradient({value_list});
+    }}
+"""
 
 AMP_LOGIC_TEMPLATE = """
     if (egr::Controller::Instance().GetCurrentAmpAttrs()->GetAmpLevel() != paddle::imperative::AmpLevel::O0){{
@@ -698,6 +704,14 @@ def _gen_out_split_and_ret_list(self, op_info, op_inst_name):
                 ret_list.append(f'{op_inst_name}.result({i})')
         return split_op_str, ret_list
 
+    def _gen_set_stop_gradient(self, ret_list):
+        if len(ret_list) > 0:
+            return SET_STOP_GRADIENT_TEMPLATE.format(
+                value_list=', '.join(ret_list)
+            )
+        else:
+            return ''
+
     def _gen_return_result(self, ret_list):
         if len(ret_list) > 1:
             return 'return std::make_tuple({});'.format(', '.join(ret_list))
@@ -1043,6 +1057,7 @@ def _gen_one_impl(
                     ),
                     set_null_type=self._gen_set_null_type(op_info, kernel_name),
                     out_split=out_split,
+                    set_stop_gradient=self._gen_set_stop_gradient(ret_list),
                     return_result=self._gen_return_result(ret_list),
                 )
                 if_inner_code = if_inner_code.split('\n')
@@ -1105,6 +1120,7 @@ def _gen_one_impl(
                 ),
                 set_null_type=self._gen_set_null_type(op_info, op_name),
                 out_split=out_split,
+                set_stop_gradient=self._gen_set_stop_gradient(ret_list),
                 return_result=self._gen_return_result(ret_list),
             )
             if op_info.is_sparse_op:
diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.cc b/paddle/fluid/pir/dialect/operator/utils/utils.cc
@@ -555,5 +555,47 @@ StringToDataLayoutMap() {
   return data_layout_map;
 }
 
+void SetStopGradient() {}
+
+void SetStopGradient(pir::Value* value) {
+  value->set_attribute(
+      "stop_gradient",
+      pir::BoolAttribute::get(pir::IrContext::Instance(), true));
+}
+
+void SetStopGradient(std::vector<pir::Value>* values) {
+  for (auto& value : *values) {
+    SetStopGradient(&value);
+  }
+}
+
+void SetStopGradient(paddle::optional<pir::Value>* value) {
+  if (value->get_ptr() != nullptr) {
+    SetStopGradient(value->get_ptr());
+  }
+}
+
+void SetStopGradient(paddle::optional<std::vector<pir::Value>>* values) {
+  if (values->get_ptr() != nullptr) {
+    SetStopGradient(values->get_ptr());
+  }
+}
+
+void PushStopGradient(const pir::Value& value, std::vector<bool>* arr) {
+  if (!IsEmptyValue(value)) {
+    arr->push_back(false);
+  } else {
+    arr->push_back(true);
+  }
+}
+
+std::vector<std::vector<bool>> ConstructStopGradient(pir::Operation* op) {
+  std::vector<std::vector<bool>> stop_gradients(op->results().size());
+  for (size_t i = 0; i < op->results().size(); i++) {
+    PushStopGradient(op->result(i), &stop_gradients[i]);
+  }
+  return stop_gradients;
+}
+
 }  // namespace dialect
 }  // namespace paddle
diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.h b/paddle/fluid/pir/dialect/operator/utils/utils.h
@@ -187,5 +187,23 @@ const std::unordered_map<std::string, phi::Place>& StringToPlaceMap();
 
 const std::unordered_map<std::string, phi::DataLayout>& StringToDataLayoutMap();
 
+void SetStopGradient();
+
+void SetStopGradient(pir::Value* value);
+
+void SetStopGradient(std::vector<pir::Value>* values);
+
+void SetStopGradient(paddle::optional<pir::Value>* value);
+
+void SetStopGradient(paddle::optional<std::vector<pir::Value>>* values);
+
+template <typename T, typename... Args>
+void SetStopGradient(T value, Args... args) {
+  SetStopGradient(&value);
+  SetStopGradient(args...);
+}
+
+std::vector<std::vector<bool>> ConstructStopGradient(pir::Operation* op);
+
 }  // namespace dialect
 }  // namespace paddle
diff --git a/paddle/fluid/primitive/codegen/templates/decomp/generated_decomp_vjp.j2 b/paddle/fluid/primitive/codegen/templates/decomp/generated_decomp_vjp.j2
@@ -134,22 +134,7 @@ std::vector<std::vector<pir::Value>> {{class_name}}::DecompVjp(pir::Operation* o
 
   VLOG(6) << "Decomp call {{fwd_name}}'s backward composite rule prepare";
 
-  std::vector<std::vector<bool>> stop_gradients(op->results().size());
-  if (op->HasAttribute(kAttrStopGradients)) {
-    auto stop_gradients_attr = op->attribute(kAttrStopGradients)
-                                   .dyn_cast<pir::ArrayAttribute>()
-                                   .AsVector();
-    {% for k in range(outputs|length) %}
-    stop_gradients[{{k}}].push_back(
-        stop_gradients_attr[{{k}}].dyn_cast<pir::BoolAttribute>().data());
-    {% endfor %}
-    VLOG(4) << " stop_gradients is set ";
-  } else {
-    {% for k in range(outputs|length) %}
-    stop_gradients[{{k}}].push_back(false);
-    {% endfor %}
-    VLOG(4) << " stop_gradients is not set ";
-  }
+  std::vector<std::vector<bool>> stop_gradients = ConstructStopGradient(op);
 
   std::vector<std::vector<paddle::Tensor>> tensor_res;
   for (auto arg : stop_gradients) {
diff --git a/paddle/fluid/pybind/imperative.cc b/paddle/fluid/pybind/imperative.cc
@@ -629,6 +629,10 @@ void BindImperative(py::module *m_ptr) {
           egr::Controller::Instance().SetCurrentTracer(tracer);
           imperative::SetCurrentTracer(tracer);
         });
+  m.def("_has_grad", []() { return egr::Controller::Instance().HasGrad(); });
+  m.def("_set_has_grad", [](bool has_grad) {
+    return egr::Controller::Instance().SetHasGrad(has_grad);
+  });
   m.def("_get_amp_attrs",
         []() { return egr::Controller::Instance().GetCurrentAmpAttrs(); });
   m.def("_set_amp_op_list",
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
@@ -1072,11 +1072,7 @@ void BindDecompVjp(pybind11::module *m) {
     for (size_t i = 0; i < decomp_res.size(); ++i) {
       py::list sub_res;
       for (size_t j = 0; j < decomp_res[i].size(); ++j) {
-        if (!decomp_res[i][j]) {
-          sub_res.append(nullptr);
-        } else {
-          sub_res.append(decomp_res[i][j]);
-        }
+        sub_res.append(decomp_res[i][j]);
       }
       res.append(sub_res);
     }
diff --git a/python/paddle/base/core.py b/python/paddle/base/core.py
@@ -298,6 +298,7 @@ def to_list(s):
         _get_phi_kernel_name,
         _get_registered_phi_kernels,
         _get_use_default_grad_op_desc_maker_ops,
+        _has_grad,
         _is_compiled_with_heterps,
         _is_dygraph_debug_enabled,
         _is_program_version_supported,
@@ -311,6 +312,7 @@ def to_list(s):
         _set_eager_deletion_mode,
         _set_fuse_parameter_group_size,
         _set_fuse_parameter_memory_size,
+        _set_has_grad,
         _set_paddle_lib_path,
         _set_warmup,
         _switch_tracer,
diff --git a/python/paddle/base/dygraph/base.py b/python/paddle/base/dygraph/base.py
@@ -306,16 +306,12 @@ def disable_dygraph() -> None:
 def _switch_tracer_mode_guard_(
     is_train: bool = True,
 ) -> Generator[None, None, None]:
-    tracer = framework._dygraph_tracer()
-    if tracer:
-        has_grad = tracer._has_grad
-        tracer._has_grad = is_train
-        try:
-            yield
-        finally:
-            tracer._has_grad = has_grad
-    else:
+    has_grad = core._has_grad()
+    core._set_has_grad(is_train)
+    try:
         yield
+    finally:
+        core._set_has_grad(has_grad)
 
 
 @overload
@@ -373,10 +369,6 @@ def no_grad(func=None):
             >>> test_layer()
 
     """
-    if in_to_static_mode():
-        warnings.warn(
-            "paddle.no_grad is only supported for inference model, and not supported for training under @to_static."
-        )
     if func is None:
         return _switch_tracer_mode_guard_(is_train=False)
     else:
@@ -455,14 +447,11 @@ def is_grad_enabled() -> bool:
             >>> paddle.is_grad_enabled()
             False
     """
-    tracer = framework._dygraph_tracer()
-    return tracer._has_grad if tracer else False
+    return core._has_grad()
 
 
 def _set_grad_enabled(mode: bool) -> None:
-    tracer = framework._dygraph_tracer()
-    if tracer:
-        tracer._has_grad = mode
+    core._set_has_grad(mode)
 
 
 class set_grad_enabled(_DecoratorContextManager):
diff --git a/test/dygraph_to_static/test_grad.py b/test/dygraph_to_static/test_grad.py
@@ -178,5 +178,27 @@ def test_run(self):
         np.testing.assert_array_equal(x.grad.numpy(), [1])
 
 
+class NoGradNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.linear = paddle.nn.Linear(3, 4)
+
+    def forward(self, x):
+        with paddle.no_grad():
+            out = self.linear(x)
+        return out
+
+
+class TestNoGrad(Dy2StTestBase):
+    @test_pir_only
+    def test_run(self):
+        net = NoGradNet()
+        net = paddle.jit.to_static(net)
+        x = paddle.rand([2, 3], 'float32')
+        x.stop_gradient = False
+        out = net(x)
+        np.testing.assert_array_equal(out.stop_gradient, True)
+
+
 if __name__ == '__main__':
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -1072,11 +1072,7 @@ void BindDecompVjp(pybind11::module *m) {`
`1072`	`1072`	`for (size_t i = 0; i < decomp_res.size(); ++i) {`
`1073`	`1073`	`py::list sub_res;`
`1074`	`1074`	`for (size_t j = 0; j < decomp_res[i].size(); ++j) {`
`1075`		`- if (!decomp_res[i][j]) {`
`1076`		`- sub_res.append(nullptr);`
`1077`		`- } else {`
`1078`		`- sub_res.append(decomp_res[i][j]);`
`1079`		`- }`
	`1075`	`+ sub_res.append(decomp_res[i][j]);`
`1080`	`1076`	`}`
`1081`	`1077`	`res.append(sub_res);`
`1082`	`1078`	`}`