[Prim][PIR] support bce_loss op forward in prim pir (#63918)

* add bce forward prim

* support dynamic

* Update composite.h

* Update test_prim_sub_graph_dynamic_shape.py

* Update composite.h

Author: Eddie-Wang1120

paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py

@@ -23,6 +23,7 @@
     "any",
     "batch_norm",
     "batch_norm_",
+    "bce_loss",
     "bmm",
     "dropout",
     "elu",
@@ -60,6 +61,7 @@
 decomp_interface_implementation_gen_op_list = [
     "any",
     "add_n",
+    "bce_loss",
     "bmm",
     "dropout",
     "elu",

paddle/fluid/primitive/composite/composite.h

@@ -251,6 +251,14 @@ Tensor reciprocal_decomp(const Tensor& x) {
   return full<T>(empty_shape, 1.0, x.dtype()) / x;
 }
 
+template <typename T>
+Tensor bce_loss_decomp(const Tensor& x, const Tensor& label) {
+  auto one = full<T>(empty_shape, 1, x.dtype());
+  auto ans = full<T>(empty_shape, -1, x.dtype()) *
+             (label * log<T>(x) + (one - label) * log<T>(one - x));
+  return ans;
+}
+
 template <typename T>
 Tensor bmm_decomp(const Tensor& x, const Tensor& y) {
   std::size_t x_ndims = x.dims().size();

test/deprecated/legacy_test/test_bce_loss.py

@@ -256,7 +256,9 @@ def setUp(self):
         self.init_test_dtype()
         self.init_test_case()
         self.op_type = "bce_loss"
+        self.prim_op_type = "comp"
         self.python_api = bce_wrapper
+        self.public_python_api = bce_wrapper
         input_np = np.random.uniform(0.1, 0.8, self.shape).astype(self.dtype)
         label_np = np.random.randint(0, 2, self.shape).astype(self.dtype)
         output_np = bce_loss(input_np, label_np)
@@ -265,7 +267,7 @@ def setUp(self):
         self.outputs = {'Out': output_np}
 
     def test_check_output(self):
-        self.check_output(check_pir=True)
+        self.check_output(check_pir=True, check_prim_pir=True)
 
     def test_check_grad(self):
         self.check_grad(['X'], 'Out', check_pir=True)
@@ -289,7 +291,7 @@ def init_test_cast(self):
 
 class TestBceLossOpFP16(TestBceLossOp):
     def test_check_output(self):
-        self.check_output(check_pir=True)
+        self.check_output(check_pir=True, check_prim_pir=True)
 
     def test_check_grad(self):
         self.check_grad(['X'], 'Out', check_pir=True)

test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py

@@ -74,6 +74,10 @@ def index_sample_net(x, index):
     return paddle.index_sample(x, index)
 
 
+def bce_loss_net(x, label):
+    return paddle._C_ops.bce_loss(x, label)
+
+
 def swiglu_net1(x, y):
     return paddle.incubate.nn.functional.swiglu(x, y)
 
@@ -318,6 +322,23 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimBceLoss(TestPrimTwo):
+    def setUp(self):
+        np.random.seed(2023)
+        self.x_shape = [20, 30, 40, 50]
+        self.y_shape = [20, 30, 40, 50]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.init_x_shape = [None, None]
+        self.init_y_shape = [None, None]
+        self.x = np.random.uniform(0.1, 0.8, self.x_shape).astype(self.dtype_x)
+        self.y = np.random.randint(0, 2, self.x_shape).astype(self.dtype_y)
+        self.net = bce_loss_net
+        self.necessary_ops = "pd_op.bce_loss"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimSwiglu1(TestPrimTwo):
     def setUp(self):
         np.random.seed(2023)