[MLU-FIX] fix(nonzero): add sync and fix its test (PaddlePaddle#391)

Author: claire199705

backends/mlu/kernels/nonzero_kernel.cc

@@ -21,8 +21,8 @@ namespace custom_kernel {
 
 template <typename T, typename Context>
 void NonZeroKernel(const Context& dev_ctx,
-                      const phi::DenseTensor& condition,
-                      phi::DenseTensor* out) {
+                   const phi::DenseTensor& condition,
+                   phi::DenseTensor* out) {
   auto dims = condition.dims();
   const int rank = dims.size();
 
@@ -38,6 +38,7 @@ void NonZeroKernel(const Context& dev_ctx,
                    GetBasePtr(&num_true));
 
   Tensor local_true_num;
+  dev_ctx.Wait();  // add sync for fully calculated results
   TensorCopy(dev_ctx, num_true, true, &local_true_num, phi::CPUPlace());
   auto true_num = *local_true_num.data<int>();
 

backends/mlu/tests/unittests/test_where_index_op_mlu.py

@@ -107,21 +107,13 @@ def test_api(self):
         paddle.set_device("CustomMLU")
         with program_guard(Program(), Program()):
             cond = paddle.static.data(name="cond", shape=[-1, 4], dtype="bool")
-            result = fluid.layers.where(cond)
+            result = paddle.nonzero(cond)
+
             exe = fluid.Executor(paddle.CustomPlace("CustomMLU", 0))
             exe.run(fluid.default_startup_program())
             cond_i = np.array([True, False, False, False]).astype("bool")
             out = exe.run(fluid.default_main_program(), feed={"cond": cond_i})
 
 
-class TestWhereRaiseError(unittest.TestCase):
-    def test_errors(self):
-        def test_type():
-            paddle.set_device("CustomMLU")
-            fluid.layers.where([10])
-
-        self.assertRaises(TypeError, test_type)
-
-
 if __name__ == "__main__":
     unittest.main()