[PIr] Fix has elementwise instruction exe bug (#60312)

zhangbo9674 · web-flow · commit 92dd1825609f · 2023-12-26T09:48:40.000+08:00
* fix

* fix
diff --git a/paddle/fluid/framework/new_executor/instruction/has_elements_instruction.cc b/paddle/fluid/framework/new_executor/instruction/has_elements_instruction.cc
@@ -56,7 +56,7 @@ HasElementsInstruction::HasElementsInstruction(
 
 void HasElementsInstruction::Run() {
   VLOG(6) << "run has_elements instruction";
-  *has_elements_ = stack_element_var_array_->empty();
+  *has_elements_ = !stack_element_var_array_->empty();
 }
 }  // namespace framework
 }  // namespace paddle
diff --git a/paddle/fluid/framework/new_executor/instruction/instruction_base.cc b/paddle/fluid/framework/new_executor/instruction/instruction_base.cc
@@ -174,6 +174,8 @@ static double GetDenseTensorEleSum(const Scope& scope,
         const phi::dtype::float16* data =
             cpu_tensor.data<phi::dtype::float16>();
         sum += static_cast<double>(data[0]);
+      } else if (cpu_tensor.dtype() == phi::DataType::BOOL) {
+        sum += static_cast<double>(cpu_tensor.data<bool>()[i]);
       } else {
         return std::numeric_limits<double>::quiet_NaN();
       }
diff --git a/paddle/fluid/framework/new_executor/instruction/tuple_pop_instruction.cc b/paddle/fluid/framework/new_executor/instruction/tuple_pop_instruction.cc
@@ -69,6 +69,7 @@ static std::stack<const Variable*> PopElements(VariableRefArray* var_array,
   for (uint64_t i = 0; i < num; i++) {
     rtn.push(var_array->back());
     var_array->pop_back();
+    VLOG(6) << "tuple pop " << rtn.top() << " from : " << var_array;
   }
   return rtn;
 }
diff --git a/paddle/fluid/framework/new_executor/instruction/tuple_push_instruction.cc b/paddle/fluid/framework/new_executor/instruction/tuple_push_instruction.cc
@@ -86,7 +86,8 @@ void TuplePushInstruction::Run() {
       DeepCopyVariable(var, copy_var, value_exe_info_, stack_size);
       VLOG(10) << "done DeepCopyVariable " << new_name;
       stack_element_var_array_->emplace_back(copy_var);
-      VLOG(6) << "push back var: " << new_name << "[" << copy_var << "]";
+      VLOG(6) << "push back var: " << new_name << "[" << copy_var << "]"
+              << "to: " << stack_element_var_array_;
     }
   }
 }
diff --git a/paddle/fluid/framework/new_executor/instruction/while_instruction.cc b/paddle/fluid/framework/new_executor/instruction/while_instruction.cc
@@ -205,6 +205,8 @@ void WhileInstruction::ShareDatasToOutputs() {
     if (out_var->IsType<phi::DenseTensor>()) {
       outputs_[i]->GetMutable<phi::DenseTensor>()->ShareDataWith(
           out_var->Get<phi::DenseTensor>());
+      VLOG(6) << "share data from " << out_var_name << "[" << out_var << "]"
+              << " -> " << i << " output[" << outputs_[i] << "]";
     } else if (out_var->IsType<phi::TensorArray>()) {
       const auto& inner_array = out_var->Get<phi::TensorArray>();
       auto* output_array = outputs_[i]->GetMutable<phi::TensorArray>();
diff --git a/test/legacy_test/test_while_loop_op.py b/test/legacy_test/test_while_loop_op.py
@@ -315,13 +315,9 @@ def body(i, x):
         np.testing.assert_allclose(np.asarray(res[0]), data, rtol=1e-05)
         np.testing.assert_allclose(np.asarray(res[1]), i_grad, rtol=1e-05)
 
-    # TODO(zhangbo): Support while grad exe for pir
-    # @test_with_pir_api
+    @test_with_pir_api
     def test_while_loop_backward2(self):
-        def cond1(i, x):
-            return i < 2
-
-        def cond2(i, x):
+        def cond(i, x):
             return i < 3
 
         def body(i, x):
@@ -339,7 +335,7 @@ def body(i, x):
             x.stop_gradient = False
             x.persistable = True
 
-            out = paddle.static.nn.while_loop(cond1, body, [i, x])
+            out = paddle.static.nn.while_loop(cond, body, [i, x])
             mean = paddle.mean(out[1])
             grad_list = append_backward(mean)
 
@@ -353,32 +349,27 @@ def body(i, x):
         feed_i = np.ones(1).astype('float32')
         feed_x = np.ones(1).astype('float32')
         data = np.asarray([2]).astype('float32')
-        ans = np.asarray([1]).astype('float32')
-        x1_grad = np.asarray([1]).astype('float32')
         i_grad = np.asarray([3]).astype('float32')
         x_grad = np.asarray([2]).astype('float32')
 
         if paddle.framework.in_pir_mode():
+            fetch_list = [out[1]]
             for p, g in grad_list:
-                if p == i:
-                    di = g
-                if p == x:
-                    dx = g
+                fetch_list.append(g)
             res = exe.run(
                 main_program,
                 feed={'i': feed_i, 'x': feed_x},
-                fetch_list=[out[1], di, dx],
+                fetch_list=fetch_list,
             )
         else:
             res = exe.run(
                 main_program,
                 feed={'i': feed_i, 'x': feed_x},
                 fetch_list=[out[1].name, i.grad_name, x.grad_name],
             )
-
-        np.testing.assert_allclose(np.asarray(res[0]), ans, rtol=1e-05)
-        np.testing.assert_allclose(np.asarray(res[1]), ans, rtol=1e-05)
-        np.testing.assert_allclose(np.asarray(res[2]), ans, rtol=1e-05)
+        np.testing.assert_allclose(np.asarray(res[0]), data, rtol=1e-05)
+        np.testing.assert_allclose(np.asarray(res[1]), i_grad, rtol=1e-05)
+        np.testing.assert_allclose(np.asarray(res[2]), x_grad, rtol=1e-05)
 
 
 class TestApiWhileLoop_NestedWithBackwardAndLoDTensorArray(unittest.TestCase):

Original file line number	Diff line number	Diff line change
`@@ -56,7 +56,7 @@ HasElementsInstruction::HasElementsInstruction(`
`56`	`56`
`57`	`57`	`void HasElementsInstruction::Run() {`
`58`	`58`	`VLOG(6) << "run has_elements instruction";`
`59`		`- *has_elements_ = stack_element_var_array_->empty();`
	`59`	`+ *has_elements_ = !stack_element_var_array_->empty();`
`60`	`60`	`}`
`61`	`61`	`} // namespace framework`
`62`	`62`	`} // namespace paddle`
Original file line number	Diff line number	Diff line change
`@@ -174,6 +174,8 @@ static double GetDenseTensorEleSum(const Scope& scope,`
`174`	`174`	`const phi::dtype::float16* data =`
`175`	`175`	`cpu_tensor.data<phi::dtype::float16>();`
`176`	`176`	`sum += static_cast<double>(data[0]);`
	`177`	`+ } else if (cpu_tensor.dtype() == phi::DataType::BOOL) {`
	`178`	`+ sum += static_cast<double>(cpu_tensor.data<bool>()[i]);`
`177`	`179`	`} else {`
`178`	`180`	`return std::numeric_limits<double>::quiet_NaN();`
`179`	`181`	`}`
Original file line number	Diff line number	Diff line change
`@@ -69,6 +69,7 @@ static std::stack<const Variable> PopElements(VariableRefArray var_array,`
`69`	`69`	`for (uint64_t i = 0; i < num; i++) {`
`70`	`70`	`rtn.push(var_array->back());`
`71`	`71`	`var_array->pop_back();`
	`72`	`+ VLOG(6) << "tuple pop " << rtn.top() << " from : " << var_array;`
`72`	`73`	`}`
`73`	`74`	`return rtn;`
`74`	`75`	`}`
Original file line number	Diff line number	Diff line change
`@@ -86,7 +86,8 @@ void TuplePushInstruction::Run() {`
`86`	`86`	`DeepCopyVariable(var, copy_var, value_exe_info_, stack_size);`
`87`	`87`	`VLOG(10) << "done DeepCopyVariable " << new_name;`
`88`	`88`	`stack_element_var_array_->emplace_back(copy_var);`
`89`		`- VLOG(6) << "push back var: " << new_name << "[" << copy_var << "]";`
	`89`	`+ VLOG(6) << "push back var: " << new_name << "[" << copy_var << "]"`
	`90`	`+ << "to: " << stack_element_var_array_;`
`90`	`91`	`}`
`91`	`92`	`}`
`92`	`93`	`}`