[Dy2St][PIR] Copy stop_gradient in while_loop fake_value and arg (#61593)

SigureMo · web-flow · commit 9882b25a6477 · 2024-02-05T11:22:22.000+08:00
diff --git a/paddle/pir/core/block_argument.cc b/paddle/pir/core/block_argument.cc
@@ -83,6 +83,11 @@ uint32_t BlockArgument::index() const {
   return IMPL_->index_;
 }
 
+const AttributeMap &BlockArgument::attributes() const {
+  CHECK_NULL_IMPL(attributes_);
+  return IMPL_->attributes_;
+}
+
 Attribute BlockArgument::attribute(const std::string &key) const {
   return impl_ ? IMPL_->attribute(key) : nullptr;
 }
diff --git a/paddle/pir/core/block_argument.h b/paddle/pir/core/block_argument.h
@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include "paddle/pir/core/operation_utils.h"
 #include "paddle/pir/core/value.h"
 namespace pir {
 class Block;
@@ -33,6 +34,7 @@ class IR_API BlockArgument : public Value {
   Block *owner() const;
   uint32_t index() const;
 
+  const AttributeMap &attributes() const;
   Attribute attribute(const std::string &key) const;
   void set_attribute(const std::string &key, Attribute value);
 
diff --git a/paddle/pir/core/region.cc b/paddle/pir/core/region.cc
@@ -52,7 +52,11 @@ void Region::CloneInto(Region &other, IrMapping &ir_mapping) const {
     auto new_block = new Block;
     ir_mapping.Add(&block, new_block);
     for (const auto &arg : block.args()) {
-      ir_mapping.Add(arg, new_block->AddArg(arg.type()));
+      auto new_arg = new_block->AddArg(arg.type());
+      ir_mapping.Add(arg, new_arg);
+      for (auto &attr : arg.dyn_cast<BlockArgument>().attributes()) {
+        new_arg.set_attribute(attr.first, attr.second);
+      }
     }
     other.push_back(new_block);
   }
diff --git a/python/paddle/static/nn/control_flow.py b/python/paddle/static/nn/control_flow.py
@@ -812,7 +812,7 @@ def create_fake_value_for_undefined_var():
             )
             cf_yield([next_cond, *unified_next_vars])
 
-            # Reset type of UndefinedVar from next_vars
+            # Reset type and stop_gradient of UndefinedVar from next_vars
             for idx, value in undefined_var_mapping.items():
                 if idx in constant_next_var_indices:
                     continue
@@ -821,6 +821,13 @@ def create_fake_value_for_undefined_var():
                 cur_block.args()[idx].set_type(value_new_type)
                 while_op.as_operation().results()[idx].set_type(value_new_type)
 
+                value_new_stop_gradient = flatten(next_vars)[idx].stop_gradient
+                value.stop_gradient = value_new_stop_gradient
+                cur_block.args()[idx].stop_gradient = value_new_stop_gradient
+                while_op.as_operation().results()[
+                    idx
+                ].stop_gradient = value_new_stop_gradient
+
         # Restore the outputs by variable and constants
         optimized_results = while_op.optimize_update()
         (optimized_variable_results,) = select_by_indices(

Original file line number	Diff line number	Diff line change
`@@ -83,6 +83,11 @@ uint32_t BlockArgument::index() const {`
`83`	`83`	`return IMPL_->index_;`
`84`	`84`	`}`
`85`	`85`
	`86`	`+const AttributeMap &BlockArgument::attributes() const {`
	`87`	`+ CHECK_NULL_IMPL(attributes_);`
	`88`	`+ return IMPL_->attributes_;`
	`89`	`+}`
	`90`	`+`
`86`	`91`	`Attribute BlockArgument::attribute(const std::string &key) const {`
`87`	`92`	`return impl_ ? IMPL_->attribute(key) : nullptr;`
`88`	`93`	`}`
Original file line number	Diff line number	Diff line change
`@@ -52,7 +52,11 @@ void Region::CloneInto(Region &other, IrMapping &ir_mapping) const {`
`52`	`52`	`auto new_block = new Block;`
`53`	`53`	`ir_mapping.Add(&block, new_block);`
`54`	`54`	`for (const auto &arg : block.args()) {`
`55`		`- ir_mapping.Add(arg, new_block->AddArg(arg.type()));`
	`55`	`+ auto new_arg = new_block->AddArg(arg.type());`
	`56`	`+ ir_mapping.Add(arg, new_arg);`
	`57`	`+ for (auto &attr : arg.dyn_cast<BlockArgument>().attributes()) {`
	`58`	`+ new_arg.set_attribute(attr.first, attr.second);`
	`59`	`+ }`
`56`	`60`	`}`
`57`	`61`	`other.push_back(new_block);`
`58`	`62`	`}`