[XPU][Phi Kernel] xpu::nonzero support simulator XPUSIM_SKIP_RUN mode

paddle/phi/kernels/xpu/masked_select_kernel.cc

@@ -14,6 +14,8 @@
 
 #include "paddle/phi/kernels/masked_select_kernel.h"
 
+#include "glog/logging.h"
+
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
 
@@ -54,7 +56,13 @@ void MaskedSelectKernel(const Context& dev_ctx,
                      mask.place(),
                      static_cast<void*>(out_size),
                      sizeof(int32_t));
-
+  if (std::getenv("XPUSIM_SKIP_RUN") &&
+      std::strcmp(std::getenv("XPUSIM_SKIP_RUN"), "1") == 0) {
+    VLOG(3) << "WARNING: In the simulator mode, the variable out_size_cpu "
+               "stores an uninitialized value. To avoid allocating a memory of "
+               "random size, we assign numel to out_size_cpu";
+    out_size_cpu = mask.numel();
+  }
   DDim out_dim{out_size_cpu};
   out->Resize(out_dim);
   auto out_data = reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(out));

paddle/phi/kernels/xpu/nonzero_kernel.cc

@@ -46,9 +46,8 @@ void NonZeroKernel(const Context& dev_ctx,
       std::strcmp(std::getenv("XPUSIM_SKIP_RUN"), "1") == 0) {
     VLOG(3) << "WARNING: In the simulator mode, the variable true_num_cpu "
                "stores an uninitialized value. To avoid allocating a memory of "
-               "random size, we limit the value of true_num_cpu to the range 0 "
-               "<= true_num_cpu < numel";
-    true_num_cpu = std::min(std::max(true_num_cpu, 0), static_cast<int>(numel));
+               "random size, we assign numel to true_num_cpu";
+    true_num_cpu = numel;
   }
 
   out->Resize(common::make_ddim({static_cast<int64_t>(true_num_cpu), rank}));

paddle/phi/kernels/xpu/sigmoid_cross_entropy_with_logits_grad_kernel.cc

@@ -16,6 +16,8 @@
 
 #include "paddle/phi/kernels/sigmoid_cross_entropy_with_logits_grad_kernel.h"
 
+#include "glog/logging.h"
+
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/backends/xpu/xpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
@@ -79,6 +81,14 @@ void SigmoidCrossEntropyWithLogitsGradKernel(
                        dev_ctx.GetPlace(),
                        static_cast<void*>(non_zero),
                        sizeof(int));
+    if (std::getenv("XPUSIM_SKIP_RUN") &&
+        std::strcmp(std::getenv("XPUSIM_SKIP_RUN"), "1") == 0) {
+      VLOG(3)
+          << "WARNING: In the simulator mode, the variable non_zero_cpu "
+             "stores an uninitialized value. To avoid allocating a memory of "
+             "random size, we assign numel to true_num_cpu";
+      non_zero_cpu = x.numel();
+    }
     r = xpu::scale(dev_ctx.x_context(),
                    reinterpret_cast<const XPUType*>(in_grad->data<T>()),
                    reinterpret_cast<XPUType*>(in_grad->data<T>()),

paddle/phi/kernels/xpu/sigmoid_cross_entropy_with_logits_kernel.cc

@@ -16,6 +16,8 @@
 
 #include "paddle/phi/kernels/sigmoid_cross_entropy_with_logits_kernel.h"
 
+#include "glog/logging.h"
+
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/backends/xpu/xpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
@@ -75,7 +77,14 @@ void SigmoidCrossEntropyWithLogitsKernel(
                        dev_ctx.GetPlace(),
                        static_cast<void*>(non_zero),
                        sizeof(int));
-
+    if (std::getenv("XPUSIM_SKIP_RUN") &&
+        std::strcmp(std::getenv("XPUSIM_SKIP_RUN"), "1") == 0) {
+      VLOG(3)
+          << "WARNING: In the simulator mode, the variable non_zero_cpu "
+             "stores an uninitialized value. To avoid allocating a memory of "
+             "random size, we assign numel to non_zero_cpu";
+      non_zero_cpu = x.numel();
+    }
     r = xpu::scale(dev_ctx.x_context(),
                    reinterpret_cast<const XPUType*>(out->data<T>()),
                    reinterpret_cast<XPUType*>(out->data<T>()),

test/xpu/test_masked_select_op_xpu.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import unittest
 
 import numpy as np
@@ -108,6 +109,20 @@ def test_static_mode(self):
         )
         self.assertEqual(np.allclose(res, np_out), True)
 
+    def test_simulator_skip_run_mode(self):
+        os.environ['XPUSIM_SKIP_RUN'] = '1'
+        paddle.disable_static(paddle.XPUPlace(0))
+        shape = (88, 6, 8)
+        np_x = np.random.random(shape).astype('float32')
+        np_mask = np.array(np.random.randint(2, size=shape, dtype=bool))
+        x = paddle.to_tensor(np_x)
+        mask = paddle.to_tensor(np_mask)
+        out = paddle.masked_select(x, mask)
+        # only check the numel of output
+        np.testing.assert_equal(out.numpy().size, np_x.size)
+        paddle.enable_static()
+        del os.environ['XPUSIM_SKIP_RUN']
+
 
 class TestMaskedSelectError(unittest.TestCase):
     def test_error(self):