[Prim][PIR] binary_cross_entropy_with_logits forward decomp

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

@@ -42,13 +42,15 @@
     "leaky_relu",
     "log_softmax",
     "mean",
+    "mean_all",
     "meshgrid",
     "one_hot",
     "p_norm",
     "pow",
     "reciprocal",
     "relu",
     "relu6",
+    "sigmoid_cross_entropy_with_logits",
     "silu",
     "swiglu",
     "softmax",
@@ -82,12 +84,14 @@
     "leaky_relu",
     "log_softmax",
     "mean",
+    "mean_all",
     "meshgrid",
     "p_norm",
     "pow",
     "reciprocal",
     "relu",
     "relu6",
+    "sigmoid_cross_entropy_with_logits",
     "silu",
     "swiglu",
     "softmax",

paddle/fluid/primitive/composite/composite.h

@@ -1161,6 +1161,73 @@ Tensor square_decomp(const Tensor& x) {
   }
 }
 
+template <typename T>
+Tensor sigmoid_cross_entropy_with_logits_decomp(
+    const Tensor& x,
+    const Tensor& label,
+    const paddle::optional<Tensor>& pos_weight,
+    bool normalize,
+    int ignore_index) {
+  auto dims = x.shape();
+  const Tensor zero = full<T>(dims, 0, x.type());
+  const Tensor one = full<T>(dims, 1, x.type());
+  Tensor pos_weight_tensor;
+  if (pos_weight) {
+    pos_weight_tensor = pos_weight.get();
+  } else {
+    pos_weight_tensor = one;
+  }
+  auto term1 = where<T>(x > zero, x, zero);
+  auto term2 = x * label;
+  auto term3 = log<T>(1 + exp<T>(-abs<T>(x)));
+  const Tensor tmp_out = term1 - term2 + term3 * pos_weight_tensor;
+  const Tensor ignore_index_tensor = full<T>(dims, ignore_index, label.type());
+  auto out = where<T>(label == ignore_index_tensor, zero, tmp_out);
+  if (normalize) {
+    // Follow the implementation in
+    // paddle/phi/kernels/cpu/sigmoid_cross_entropy_with_logits_kernel.cc
+    const Tensor eps1 = full<T>(dims, 1e-6, x.type());
+    auto diff = label - ignore_index_tensor;
+    const Tensor tmp_norm = sum<T>(where<T>(abs<T>(diff) > eps1, one, zero));
+    // Follow the implementation in
+    // paddle/phi/kernels/cpu/sigmoid_cross_entropy_with_logits_kernel.cc
+    const Tensor eps2 = full<T>(empty_shape, 1e-5, x.type());
+    auto norm = where<T>(tmp_norm > eps2, tmp_norm, eps2);
+    out = out / norm;
+  }
+  return out;
+}
+
+template <typename T>
+Tensor mean_all_decomp(const Tensor& x) {
+  auto org_dtype = x.dtype();
+  auto x_cast = x;
+  auto x_shape = x.shape();
+  bool need_cast = is_half_dtype(org_dtype);
+  if (need_cast) {
+    x_cast = cast<T>(x, DataType::FLOAT32);
+  }
+
+  Tensor ans;
+  if (has_dynamic_shape(x_shape)) {
+    Tensor x_shape_tensor = shape<T>(x_cast);
+    Tensor value = get_slice<T>(x_shape_tensor, 0);
+    for (size_t i = 1; i < x_shape.size(); i++) {
+      value = value * get_slice<T>(x_shape_tensor, i);
+    }
+    value = reshape<T>(value, {});
+    ans = sum<T>(x_cast) / cast<T>(value, x_cast.dtype());
+  } else {
+    ans = sum<T>(x_cast) / x_cast.numel();
+  }
+
+  if (need_cast) {
+    return cast<T>(ans, org_dtype);
+  } else {
+    return ans;
+  }
+}
+
 template <typename T>
 Tensor embedding_decomp(const Tensor& x,
                         const Tensor& weight,

test/deprecated/legacy_test/test_sigmoid_cross_entropy_with_logits_op.py

@@ -38,6 +38,8 @@ class TestSigmoidCrossEntropyWithLogitsOp1(OpTest):
     def setUp(self):
         self.op_type = "sigmoid_cross_entropy_with_logits"
         self.python_api = loss_wrapper
+        self.prim_op_type = "comp"
+        self.public_python_api = loss_wrapper
         batch_size = 64
         num_classes = 20
         self.inputs = {
@@ -60,10 +62,10 @@ def setUp(self):
         self.outputs = {'Out': -term1 - term2}
 
     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)
+        self.check_grad(['X'], 'Out', check_pir=True, check_prim_pir=True)
 
 
 class TestSigmoidCrossEntropyWithLogitsOp2(OpTest):
@@ -72,6 +74,8 @@ class TestSigmoidCrossEntropyWithLogitsOp2(OpTest):
     def setUp(self):
         self.op_type = "sigmoid_cross_entropy_with_logits"
         self.python_api = loss_wrapper
+        self.prim_op_type = "comp"
+        self.public_python_api = loss_wrapper
         batch_size = 64
         num_classes = 20
         ignore_index = -1
@@ -99,10 +103,10 @@ def setUp(self):
         self.outputs = {'Out': out}
 
     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)
+        self.check_grad(['X'], 'Out', check_pir=True, check_prim_pir=True)
 
 
 class TestSigmoidCrossEntropyWithLogitsOp3(OpTest):
@@ -111,6 +115,8 @@ class TestSigmoidCrossEntropyWithLogitsOp3(OpTest):
     def setUp(self):
         self.op_type = "sigmoid_cross_entropy_with_logits"
         self.python_api = loss_wrapper
+        self.prim_op_type = "comp"
+        self.public_python_api = loss_wrapper
         batch_size = 64
         num_classes = 20
         self.inputs = {
@@ -133,10 +139,10 @@ def setUp(self):
         self.outputs = {'Out': -term1 - term2}
 
     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)
+        self.check_grad(['X'], 'Out', check_pir=True, check_prim_pir=True)
 
 
 class TestSigmoidCrossEntropyWithLogitsOp4(OpTest):
@@ -145,6 +151,8 @@ class TestSigmoidCrossEntropyWithLogitsOp4(OpTest):
     def setUp(self):
         self.op_type = "sigmoid_cross_entropy_with_logits"
         self.python_api = loss_wrapper
+        self.prim_op_type = "comp"
+        self.public_python_api = loss_wrapper
         batch_size = 64
         num_classes = 20
 
@@ -171,7 +179,7 @@ def setUp(self):
         self.outputs = {'Out': term1 - term2 + term3}
 
     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', max_relative_error=0.0005, check_pir=True)
@@ -181,6 +189,8 @@ class TestSigmoidCrossEntropyWithNorm(OpTest):
     def setUp(self):
         self.op_type = "sigmoid_cross_entropy_with_logits"
         self.python_api = loss_wrapper
+        self.prim_op_type = "comp"
+        self.public_python_api = loss_wrapper
         batch_size = 64
         num_classes = 20
         ignore_index = -1
@@ -207,10 +217,10 @@ def setUp(self):
         self.outputs = {'Out': out}
 
     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)
+        self.check_grad(['X'], 'Out', check_pir=True, check_prim_pir=True)
 
 
 class TestSigmoidCrossEntropyWithLogitsOp5(OpTest):
@@ -219,6 +229,8 @@ class TestSigmoidCrossEntropyWithLogitsOp5(OpTest):
     def setUp(self):
         self.op_type = "sigmoid_cross_entropy_with_logits"
         self.python_api = loss_wrapper
+        self.prim_op_type = "comp"
+        self.public_python_api = loss_wrapper
         batch_size = [10, 10]
         num_classes = 20
         self.inputs = {
@@ -241,16 +253,18 @@ def setUp(self):
         self.outputs = {'Out': -term1 - term2}
 
     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)
+        self.check_grad(['X'], 'Out', check_pir=True, check_prim_pir=True)
 
 
 class TestSigmoidCrossEntropyWithNorm2(OpTest):
     def setUp(self):
         self.op_type = "sigmoid_cross_entropy_with_logits"
         self.python_api = loss_wrapper
+        self.prim_op_type = "comp"
+        self.public_python_api = loss_wrapper
         batch_size = [10, 10]
         num_classes = 20
         ignore_index = -1
@@ -277,10 +291,10 @@ def setUp(self):
         self.outputs = {'Out': out}
 
     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)
+        self.check_grad(['X'], 'Out', check_pir=True, check_prim_pir=True)
 
 
 class TestSigmoidCrossEntropyWithLogitsOp6(OpTest):
@@ -289,6 +303,8 @@ class TestSigmoidCrossEntropyWithLogitsOp6(OpTest):
     def setUp(self):
         self.op_type = "sigmoid_cross_entropy_with_logits"
         self.python_api = loss_wrapper
+        self.prim_op_type = "comp"
+        self.public_python_api = loss_wrapper
         batch_size = [10, 10]
         num_classes = 20
         self.inputs = {
@@ -311,10 +327,10 @@ def setUp(self):
         self.outputs = {'Out': -term1 - term2}
 
     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)
+        self.check_grad(['X'], 'Out', check_pir=True, check_prim_pir=True)
 
 
 class TestSigmoidCrossEntropyWithLogitsOpError(unittest.TestCase):

test/legacy_test/test_binary_cross_entropy_with_logits_op.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+
+
+class TestBinaryCrossEntropyWithLogits(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.x = np.random.randn(300, 1000).astype("float32")
+        self.y = np.random.randint(0, 2, (300, 1000)).astype("float32")
+        self.logits = paddle.to_tensor(self.x)
+        self.labels = paddle.to_tensor(self.y)
+        self.weight = paddle.to_tensor(
+            np.random.randn(300, 1000).astype("float32")
+        )
+        self.reduction = ["none", "mean", "sum"]
+        self.pos_weight = paddle.to_tensor(
+            np.random.randn(1000).astype("float32")
+        )
+
+    def test_binary_cross_entropy_with_logits(self):
+        for reduction in self.reduction:
+            dynamic_result = (
+                paddle.nn.functional.binary_cross_entropy_with_logits(
+                    self.logits,
+                    self.labels,
+                    weight=self.weight,
+                    reduction=reduction,
+                    pos_weight=self.pos_weight,
+                )
+            )
+            paddle.core._set_prim_all_enabled(True)
+            static_result = paddle.jit.to_static(
+                paddle.nn.functional.binary_cross_entropy_with_logits,
+                full_graph=True,
+            )(
+                self.logits,
+                self.labels,
+                weight=self.weight,
+                reduction=reduction,
+                pos_weight=self.pos_weight,
+            )
+            paddle.core._set_prim_all_enabled(False)
+            np.testing.assert_allclose(
+                dynamic_result.numpy(), static_result.numpy(), rtol=1e-4
+            )
+
+
+class TestBinaryCrossEntropyWithLogits1(TestBinaryCrossEntropyWithLogits):
+    def setUp(self):
+        super().setUp()
+        self.weight = None
+
+
+class TestBinaryCrossEntropyWithLogits2(TestBinaryCrossEntropyWithLogits):
+    def setUp(self):
+        super().setUp()
+        self.pos_weight = None
+
+
+class TestBinaryCrossEntropyWithLogits3(TestBinaryCrossEntropyWithLogits):
+    def setUp(self):
+        super().setUp()
+        self.weight = None
+        self.pos_weight = None
+
+
+if __name__ == "__main__":
+    unittest.main()

test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py

@@ -90,6 +90,10 @@ def dropout_net1(x):
     return paddle.nn.functional.dropout(x, 0.5)
 
 
+def mean_all_net1(x):
+    return paddle._C_ops.mean_all(x)
+
+
 group_norm1 = paddle.nn.GroupNorm(num_channels=128, num_groups=32)
 
 
@@ -509,5 +513,18 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimMeanAll(TestPrimBase):
+    def setUp(self):
+        np.random.seed(2023)
+        paddle.seed(2023)
+        self.shape_x = [300, 4096]
+        self.dtype_x = "float32"
+        self.init_x_shape = [None, 4096]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.net = mean_all_net1
+        self.necessary_ops = "pd_op.mean_all"
+        self.enable_cinn = False
+
+
 if __name__ == "__main__":
     unittest.main()