[PIR] support while api for python api. (#60364)

winter-wang · web-flow · commit 7d9f2a2c9efc · 2023-12-28T10:55:20.000+08:00
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -317,11 +317,15 @@ void BindBlock(py::module *m) {
         The constructor of Block should not be invoked directly. You can
         use `Program.block()` to get a block.
   )DOC");
-  block
+  block.def("empty", &Block::empty)
       .def(
           "front",
           [](Block &self) { return &self.front(); },
           return_value_policy::reference)
+      .def(
+          "back",
+          [](Block &self) { return &self.back(); },
+          return_value_policy::reference)
       .def_property_readonly(
           "parent_op",
           [](Block &self) { return self.GetParentOp(); },
diff --git a/paddle/pir/core/block.cc b/paddle/pir/core/block.cc
@@ -73,10 +73,7 @@ Operation *Block::Take(Operation *op) {
   return op;
 }
 
-void Block::SetParent(Region *parent, Region::Iterator position) {
-  parent_ = parent;
-  position_ = position;
-}
+void Block::SetParent(Region *parent) { parent_ = parent; }
 
 Block::UseIterator Block::use_begin() const { return first_use_; }
 
diff --git a/paddle/pir/core/block.h b/paddle/pir/core/block.h
@@ -73,7 +73,6 @@ class IR_API Block {
   Iterator insert(ConstIterator iterator, Operation *op);
   Iterator erase(ConstIterator position);
   void clear();
-  operator Region::Iterator() { return position_; }
 
   // Assign the operation underlying in position with parameter op,
   // meanwhile, destroy the original operation.
@@ -145,7 +144,7 @@ class IR_API Block {
 
   // Allow access to 'SetParent'.
   friend class Region;
-  void SetParent(Region *parent, Region::Iterator position);
+  void SetParent(Region *parent);
 
   // Take out corresponding Operation and its ownershipe.
   friend class Operation;
@@ -154,7 +153,6 @@ class IR_API Block {
   static bool TopoOrderCheck(const OpListType &op_list);
 
  private:
-  Region::Iterator position_;
   BlockOperand first_use_;
   OpListType ops_;         // owned
   ArgListType arguments_;  // owned
diff --git a/paddle/pir/core/region.cc b/paddle/pir/core/region.cc
@@ -32,7 +32,7 @@ void Region::push_front(Block *block) { insert(blocks_.begin(), block); }
 
 Region::Iterator Region::insert(ConstIterator position, Block *block) {
   Region::Iterator iter = blocks_.insert(position, block);
-  block->SetParent(this, iter);
+  block->SetParent(this);
   return iter;
 }
 
@@ -54,7 +54,7 @@ void Region::TakeBody(Region &&other) {
   clear();
   blocks_.swap(other.blocks_);
   for (auto iter = blocks_.begin(); iter != blocks_.end(); ++iter) {
-    (*iter)->SetParent(this, iter);
+    (*iter)->SetParent(this);
   }
 }
 
@@ -72,11 +72,11 @@ void Region::clear() {
 
 void Region::swap(Region &&other) {
   blocks_.swap(other.blocks_);
-  for (auto iter = begin(); iter != end(); ++iter) {
-    iter->SetParent(this, iter);
+  for (auto &block : *this) {
+    block.SetParent(this);
   }
-  for (auto iter = other.begin(); iter != other.end(); ++iter) {
-    iter->SetParent(&other, iter);
+  for (auto &block : other) {
+    block.SetParent(&other);
   }
 }
 
diff --git a/python/paddle/static/nn/control_flow.py b/python/paddle/static/nn/control_flow.py
@@ -153,14 +153,21 @@ class WhileGuard(BlockGuard):
     def __init__(self, while_op):
         if not isinstance(while_op, While):
             raise TypeError("WhileGuard takes a while op")
-        super().__init__(while_op.helper.main_program)
+        if not in_pir_mode():
+            super().__init__(while_op.helper.main_program)
         self.while_op = while_op
 
     def __enter__(self):
+        if in_pir_mode():
+            self.block = build_while_op(self.while_op.cond_var, []).body()
+            return self.block.__enter__()
         self.while_op.status = While.IN_WHILE_BLOCK
         return super().__enter__()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        if in_pir_mode():
+            cf_yield([self.while_op.cond_var])
+            return self.block.__exit__(exc_type, exc_val, exc_tb)
         if exc_type is not None:
             return False
         self.while_op.status = While.AFTER_WHILE_BLOCK
@@ -509,16 +516,18 @@ class While:
     AFTER_WHILE_BLOCK = 2
 
     def __init__(self, cond, is_test=False, name=None):
-        self.helper = LayerHelper("while", name=name)
-        self.status = While.BEFORE_WHILE_BLOCK
+        self.cond_var = cond
         check_variable_and_dtype(cond, 'cond', ['bool'], 'static.nn.While')
         if reduce(lambda a, b: a * b, cond.shape, 1) != 1:
             raise TypeError(
                 "condition expected shape as [1], but given shape as {}.".format(
                     list(cond.shape)
                 )
             )
-        self.cond_var = cond
+        if in_pir_mode():
+            return
+        self.status = While.BEFORE_WHILE_BLOCK
+        self.helper = LayerHelper("while", name=name)
         self.is_test = is_test
 
     def block(self):
@@ -1870,5 +1879,4 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         self.inside_scope = False
         if exc_type is not None:
             return False  # re-raise exception
-
         return True
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -4719,12 +4719,15 @@ def increment(x, value=1.0, name=None):
             [1.])
 
     """
-    if in_dynamic_or_pir_mode():
+    if in_dynamic_mode():
+        return _C_ops.increment_(x, value)
+
+    check_variable_and_dtype(
+        x, 'x', ['float32', 'float64', 'int32', 'int64'], 'increment'
+    )
+    if in_pir_mode():
         return _C_ops.increment_(x, value)
     else:
-        check_variable_and_dtype(
-            x, 'x', ['float32', 'float64', 'int32', 'int64'], 'increment'
-        )
         helper = LayerHelper("increment", **locals())
         helper.append_op(
             type='increment',
diff --git a/test/legacy_test/test_while_op.py b/test/legacy_test/test_while_op.py
@@ -12,17 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import unittest
 
 import numpy
 from utils import compare_legacy_with_pt
 
 import paddle
-from paddle import base
+from paddle import base, set_flags
 from paddle.base import core
 from paddle.base.backward import append_backward
 from paddle.base.executor import Executor
+from paddle.base.framework import in_pir_mode
 from paddle.incubate.layers.nn import shuffle_batch
+from paddle.pir_utils import test_with_pir_api
 
 paddle.enable_static()
 
@@ -70,7 +73,7 @@ def simple_net(self):
                 prev2 = paddle.tensor.array_read(array=mem_array, i=j)
                 result2 = paddle.add_n([d2, prev2])
 
-                j = paddle.increment(x=j)
+                paddle.increment(x=j)
                 paddle.tensor.array_write(result2, i=j, array=mem_array)
                 paddle.assign(paddle.less_than(x=j, y=array_len2), cond2)
 
@@ -79,7 +82,8 @@ def simple_net(self):
         loss = paddle.mean(sum_result)
         return loss, sum_result
 
-    # TODO(zhangbo): Support pir test(support write_to_array and read_from_array, support while_grad).
+    # TODO(winter-wang): Support pir test in (FLAGS_enable_pir_in_executor_trace_run = False && FLAGS_new_executor_serial_run == False).
+    @test_with_pir_api
     def test_simple_net(self):
         main_program = base.Program()
         startup_program = base.Program()
@@ -88,6 +92,14 @@ def test_simple_net(self):
 
             append_backward(loss)
 
+            if in_pir_mode():
+                flag_1 = "FLAGS_enable_pir_in_executor_trace_run"
+                flag_2 = "FLAGS_new_executor_serial_run"
+                os.environ[flag_1] = 'True'
+                os.environ[flag_2] = 'True'
+                set_flags({flag_1: True})
+                set_flags({flag_2: True})
+
             cpu = core.CPUPlace()
             exe = Executor(cpu)
             d = []
@@ -99,26 +111,48 @@ def test_simple_net(self):
                 feed={'d0': d[0], 'd1': d[1], 'd2': d[2]},
                 fetch_list=[sum_result],
             )
+            if in_pir_mode():
+                del os.environ[flag_1]
+                del os.environ[flag_2]
+                set_flags({flag_1: False})
+                set_flags({flag_2: False})
             self.assertAlmostEqual(numpy.sum(d), numpy.sum(outs[0]), delta=0.01)
 
-    # TODO(zhangbo): Support pir test(support write_to_array and read_from_array)
+    # TODO(winter-wang): Support pir test in (FLAGS_enable_pir_in_executor_trace_run = False && FLAGS_new_executor_serial_run == False).
+    @test_with_pir_api
     def test_simple_net_forward(self):
         main_program = base.Program()
         startup_program = base.Program()
         with base.program_guard(main_program, startup_program):
             self.simple_net()
-            binary = base.compiler.CompiledProgram(main_program)
+            if in_pir_mode():
+                binary = main_program
+            else:
+                binary = base.compiler.CompiledProgram(main_program)
             cpu = core.CPUPlace()
             exe = Executor(cpu)
             d = []
 
             for i in range(3):
                 d.append(numpy.random.random(size=[10]).astype('float32'))
 
+            if in_pir_mode():
+                flag_1 = "FLAGS_enable_pir_in_executor_trace_run"
+                flag_2 = "FLAGS_new_executor_serial_run"
+                os.environ[flag_1] = 'True'
+                os.environ[flag_2] = 'True'
+                set_flags({flag_1: True})
+                set_flags({flag_2: True})
             for _ in range(2):
                 exe.run(binary, feed={'d0': d[0], 'd1': d[1], 'd2': d[2]})
+            if in_pir_mode():
+                del os.environ[flag_1]
+                del os.environ[flag_2]
+                set_flags({flag_1: False})
+                set_flags({flag_2: False})
 
     @compare_legacy_with_pt
+    @test_with_pir_api
     def test_exceptions(self):
         i = paddle.zeros(shape=[2], dtype='int64')
         array_len = paddle.tensor.fill_constant(
@@ -134,6 +168,7 @@ def test_exceptions(self):
 
 class BadInputTest(unittest.TestCase):
     @compare_legacy_with_pt
+    @test_with_pir_api
     def test_error(self):
         with base.program_guard(base.Program()):
 
@@ -158,8 +193,9 @@ def body_func(i, ten, batch_info, origin_seq):
 
         x = paddle.static.data(name='x', shape=[-1, 1, 4], dtype='float32')
         y = paddle.static.data(name='y', shape=[-1, 1, 1], dtype='float32')
-        x.desc.set_need_check_feed(False)
-        y.desc.set_need_check_feed(False)
+        if not in_pir_mode():
+            x.desc.set_need_check_feed(False)
+            y.desc.set_need_check_feed(False)
         temp = paddle.concat([x, y], axis=-1)
 
         i = paddle.tensor.fill_constant(shape=[1], value=0, dtype='int32')
@@ -190,6 +226,7 @@ def body_func(i, ten, batch_info, origin_seq):
 
 class TestOutputsMustExistsInputs(unittest.TestCase):
     @compare_legacy_with_pt
+    @test_with_pir_api
     def test_outputs_exists_inputs(self):
         """
         We guarantee that the output tensor must be in the input tensor, so that the output and input can correspond to each other, but the input can be greater than the number of outputs. It's required in paddle2onnx.
@@ -218,17 +255,20 @@ def body(i, s, x):
             paddle.enable_static()
             x = paddle.static.data(shape=[-1], name='x', dtype='float32')
             func(x)
-        for op in main_program.block(0).ops:
-            if op.type == "while":
-                for out_name in op.output("Out"):
-                    if out_name in op.input("Condition"):
-                        continue
-                    self.assertTrue(
-                        out_name in op.input("X"),
-                        "In while op, the variable in output(`Out`) must exists in inputs(`X`), but the variable with name `{}` not meet the precondition.".format(
-                            out_name
-                        ),
-                    )
+
+        # NOTE(winter-wang): The while_op in pir mode  doesn't need following constrait, so hre only check when in non-pir mode.
+        if not in_pir_mode():
+            for op in main_program.block(0).ops:
+                if op.type == "while":
+                    for out_name in op.output("Out"):
+                        if out_name in op.input("Condition"):
+                            continue
+                        self.assertTrue(
+                            out_name in op.input("X"),
+                            "In while op, the variable in output(`Out`) must exists in inputs(`X`), but the variable with name `{}` not meet the precondition.".format(
+                                out_name
+                            ),
+                        )
 
 
 if __name__ == '__main__':

Original file line number	Diff line number	Diff line change
`@@ -73,10 +73,7 @@ Operation Block::Take(Operation op) {`
`73`	`73`	`return op;`
`74`	`74`	`}`
`75`	`75`
`76`		`-void Block::SetParent(Region *parent, Region::Iterator position) {`
`77`		`- parent_ = parent;`
`78`		`- position_ = position;`
`79`		`-}`
	`76`	`+void Block::SetParent(Region *parent) { parent_ = parent; }`
`80`	`77`
`81`	`78`	`Block::UseIterator Block::use_begin() const { return first_use_; }`
`82`	`79`
Original file line number	Diff line number	Diff line change
`@@ -32,7 +32,7 @@ void Region::push_front(Block *block) { insert(blocks_.begin(), block); }`
`32`	`32`
`33`	`33`	`Region::Iterator Region::insert(ConstIterator position, Block *block) {`
`34`	`34`	`Region::Iterator iter = blocks_.insert(position, block);`
`35`		`- block->SetParent(this, iter);`
	`35`	`+ block->SetParent(this);`
`36`	`36`	`return iter;`
`37`	`37`	`}`
`38`	`38`
`@@ -54,7 +54,7 @@ void Region::TakeBody(Region &&other) {`
`54`	`54`	`clear();`
`55`	`55`	`blocks_.swap(other.blocks_);`
`56`	`56`	`for (auto iter = blocks_.begin(); iter != blocks_.end(); ++iter) {`
`57`		`- (*iter)->SetParent(this, iter);`
	`57`	`+ (*iter)->SetParent(this);`
`58`	`58`	`}`
`59`	`59`	`}`
`60`	`60`
`@@ -72,11 +72,11 @@ void Region::clear() {`
`72`	`72`
`73`	`73`	`void Region::swap(Region &&other) {`
`74`	`74`	`blocks_.swap(other.blocks_);`
`75`		`- for (auto iter = begin(); iter != end(); ++iter) {`
`76`		`- iter->SetParent(this, iter);`
	`75`	`+ for (auto &block : *this) {`
	`76`	`+ block.SetParent(this);`
`77`	`77`	`}`
`78`		`- for (auto iter = other.begin(); iter != other.end(); ++iter) {`
`79`		`- iter->SetParent(&other, iter);`
	`78`	`+ for (auto &block : other) {`
	`79`	`+ block.SetParent(&other);`
`80`	`80`	`}`
`81`	`81`	`}`
`82`	`82`