[Prim][PIR] support huber_loss op forward in prim pir (#64425)

* update huber_loss

* update dynamic

Author: Eddie-Wang1120

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

@@ -56,6 +56,7 @@
     "squeeze",
     "stack",
     "unsqueeze",
+    "huber_loss",
 ]
 
 # come into effect in generated file op_decomp.cc
@@ -95,6 +96,7 @@
     "squeeze",
     "stack",
     "unsqueeze",
+    "huber_loss",
 ]
 
 # xshape output will no longer used after decomp, but return none to keep output num the same as origin op

paddle/fluid/primitive/composite/composite.h

@@ -208,6 +208,25 @@ Tensor pow_decomp(const Tensor& x, const paddle::Scalar& y) {
   }
 }
 
+template <typename T>
+std::tuple<Tensor, Tensor> huber_loss_decomp(const Tensor& input,
+                                             const Tensor& label,
+                                             float delta) {
+  Tensor delta_full;
+  if (has_dynamic_shape(input.shape())) {
+    delta_full =
+        backend::full_with_tensor<T>(shape<T>(input), delta, input.dtype());
+  } else {
+    delta_full = full<T>(input.shape(), delta, input.dtype());
+  }
+  auto val = label - input;
+  auto abs_val = abs<T>(val);
+  auto ans = where<T>(abs_val <= delta_full,
+                      0.5 * val * val,
+                      delta_full * (abs_val - 0.5 * delta_full));
+  return std::make_tuple(ans, val);
+}
+
 template <typename T>
 Tensor one_hot_decomp(const Tensor& x, const Tensor& num_classes) {
   auto num_classes_tensor =

paddle/phi/infermeta/binary.cc

@@ -2069,6 +2069,8 @@ void HuberLossInferMeta(const MetaTensor& input,
   auto out_dims = label_dims;
   residual->set_dims(out_dims);
   out->set_dims(out_dims);
+  out->set_dtype(input.dtype());
+  residual->set_dtype(input.dtype());
   out->share_lod(input);
 }
 

test/deprecated/legacy_test/test_huber_loss_op.py

@@ -29,16 +29,18 @@ def huber_loss_forward(val, delta):
         return delta * (abs_val - 0.5 * delta)
 
 
-def huber_loss_wraper(x, y, delta):
+def huber_loss_wrapper(x, y, delta):
     a = paddle._C_ops.huber_loss(x, y, delta)
     return a
 
 
 class TestHuberLossOp(OpTest):
     def setUp(self):
         self.op_type = 'huber_loss'
+        self.prim_op_type = "comp"
         self.python_out_sig = ["Out"]
-        self.python_api = huber_loss_wraper
+        self.python_api = huber_loss_wrapper
+        self.public_python_api = huber_loss_wrapper
 
         self.delta = 1.0
         self.init_dtype()
@@ -65,7 +67,7 @@ def set_shape(self):
         return (100, 1)
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_prim_pir=True)
 
     def test_check_grad_normal(self):
         self.check_grad(['X', 'Y'], 'Out')
@@ -105,8 +107,10 @@ def init_dtype(self):
 class TestHuberLossBF16Op(OpTest):
     def setUp(self):
         self.op_type = 'huber_loss'
+        self.prim_op_type = "comp"
         self.python_out_sig = ["Out"]
-        self.python_api = huber_loss_wraper
+        self.python_api = huber_loss_wrapper
+        self.public_python_api = huber_loss_wrapper
 
         self.delta = 1.0
         self.init_dtype()
@@ -142,7 +146,7 @@ def set_shape(self):
         return (100, 1)
 
     def test_check_output(self):
-        self.check_output_with_place(self.place)
+        self.check_output_with_place(self.place, check_prim_pir=True)
 
     def test_check_grad_normal(self):
         self.check_grad_with_place(self.place, ['X', 'Y'], 'Out')

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 huber_loss_net(x, label):
+    return paddle._C_ops.huber_loss(x, label, 1.0)
+
+
 def bce_loss_net(x, label):
     return paddle._C_ops.bce_loss(x, label)
 
@@ -322,6 +326,23 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimHuberLoss(TestPrimTwo):
+    def setUp(self):
+        np.random.seed(2023)
+        self.x_shape = [100, 1]
+        self.y_shape = [100, 1]
+        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, self.x_shape).astype(self.dtype_x)
+        self.y = np.random.uniform(0, 1.0, self.y_shape).astype(self.dtype_y)
+        self.net = huber_loss_net
+        self.necessary_ops = "pd_op.huber_loss"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimBceLoss(TestPrimTwo):
     def setUp(self):
         np.random.seed(2023)