[Prim][PIR] Forward decomposite the lerp op (#65967)

zeroRains · web-flow · commit 15ef3cbc4d02 · 2024-07-16T16:41:37.000+08:00
* forward decomposite the lerp op

* lerp

* fix the bug in the get_output_dims

* polish

* fix code style

* move the modify to infermeta

* fix the bug

* fix the bug
diff --git a/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py b/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py
@@ -41,6 +41,7 @@
     "instance_norm",
     "layer_norm",
     "leaky_relu",
+    "lerp",
     "log_loss",
     "log_softmax",
     "mean",
@@ -87,6 +88,7 @@
     "instance_norm",
     "layer_norm",
     "leaky_relu",
+    "lerp",
     "log_loss",
     "log_softmax",
     "mean",
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -1514,6 +1514,32 @@ Tensor elu_decomp(const Tensor& x, const float alpha) {
   }
 }
 
+template <typename T>
+Tensor lerp_decomp(const Tensor& x, const Tensor& y, const Tensor& weight) {
+  Tensor x_cast = x;
+  Tensor y_cast = y;
+  Tensor weight_cast = weight;
+  bool need_cast = false;
+  if (is_half_dtype(x.dtype())) {
+    need_cast = true;
+    x_cast = cast<T>(x, DataType::FLOAT32);
+  }
+  if (is_half_dtype(y.dtype())) {
+    need_cast = true;
+    y_cast = cast<T>(y, DataType::FLOAT32);
+  }
+  if (is_half_dtype(weight.dtype())) {
+    need_cast = true;
+    weight_cast = cast<T>(weight, DataType::FLOAT32);
+  }
+  Tensor res = x_cast + weight_cast * (y_cast - x_cast);
+  if (need_cast) {
+    return cast<T>(res, x.dtype());
+  } else {
+    return res;
+  }
+}
+
 template <typename T>
 Tensor log_loss_decomp(const Tensor& input,
                        const Tensor& label,
diff --git a/paddle/phi/infermeta/ternary.cc b/paddle/phi/infermeta/ternary.cc
@@ -1142,9 +1142,66 @@ void LerpInferMeta(const MetaTensor& x,
   auto x_dims = x.dims();
   auto y_dims = y.dims();
   auto w_dims = weight.dims();
-  DDim out_dims;
-  out_dims = funcs::GetOutputDims(x_dims, y_dims);
-  out_dims = funcs::GetOutputDims(out_dims, w_dims);
+  DDim l_dims, s_dims;
+  if (x_dims.size() > y_dims.size()) {
+    l_dims = x_dims;
+    s_dims = y_dims;
+  } else {
+    l_dims = y_dims;
+    s_dims = x_dims;
+  }
+  std::vector<int64_t> shapes = common::vectorize<int64_t>(l_dims);
+  for (int i = s_dims.size() - 1, j = l_dims.size() - 1; i >= 0; --i, --j) {
+    int64_t s = s_dims[i];
+    int64_t l = l_dims[j];
+    if (s != l) {
+      if (l == 1) {
+        shapes[j] = s;
+      } else if (s == 1 || s == -1) {
+        shapes[j] = l;
+      } else if (l == -1) {
+        shapes[j] = s;
+      } else {
+        PADDLE_THROW(errors::InvalidArgument(
+            "The shape of tensor a %s:%d must match shape of tensor b "
+            "%s:%d.",
+            s_dims.to_str(),
+            i,
+            l_dims.to_str(),
+            j));
+      }
+    }
+  }
+  if (static_cast<int>(shapes.size()) > w_dims.size()) {
+    l_dims = common::make_ddim(shapes);
+    s_dims = w_dims;
+  } else {
+    l_dims = w_dims;
+    s_dims = common::make_ddim(shapes);
+  }
+  std::vector<int64_t> shapes_out = common::vectorize<int64_t>(l_dims);
+  for (int i = s_dims.size() - 1, j = l_dims.size() - 1; i >= 0; --i, --j) {
+    int64_t s = s_dims[i];
+    int64_t l = l_dims[j];
+    if (s != l) {
+      if (l == 1) {
+        shapes_out[j] = s;
+      } else if (s == 1 || s == -1) {
+        shapes_out[j] = l;
+      } else if (l == -1) {
+        shapes_out[j] = s;
+      } else {
+        PADDLE_THROW(errors::InvalidArgument(
+            "The shape of tensor a %s:%d must match shape of tensor b "
+            "%s:%d.",
+            s_dims.to_str(),
+            i,
+            l_dims.to_str(),
+            j));
+      }
+    }
+  }
+  DDim out_dims = common::make_ddim(shapes_out);
   out->set_dims(out_dims);
   out->set_dtype(x.dtype());
   out->share_lod(x);
diff --git a/test/legacy_test/test_lerp_op.py b/test/legacy_test/test_lerp_op.py
@@ -29,6 +29,8 @@ class TestLerp(OpTest):
     def setUp(self):
         self.op_type = "lerp"
         self.python_api = paddle.lerp
+        self.prim_op_type = "comp"
+        self.public_python_api = paddle.lerp
         self.init_dtype()
         self.init_shape()
         self.init_xyshape()
@@ -53,10 +55,10 @@ def init_wshape(self):
         self.wshape = [1]
 
     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', 'Y'], 'Out', check_pir=True)
+        self.check_grad(['X', 'Y'], 'Out', check_pir=True, check_prim_pir=True)
 
 
 class TestLerpWithDim2(TestLerp):
@@ -231,6 +233,8 @@ class TestLerpBF16(TestLerp):
     def setUp(self):
         self.op_type = "lerp"
         self.python_api = paddle.lerp
+        self.prim_op_type = "comp"
+        self.public_python_api = paddle.lerp
         self.dtype = np.uint16
         self.init_shape()
         self.init_xyshape()
@@ -270,7 +274,7 @@ def init_grad(self, w):
 
     def test_check_output(self):
         place = core.CUDAPlace(0)
-        self.check_output_with_place(place, check_pir=True)
+        self.check_output_with_place(place, check_pir=True, check_prim_pir=True)
 
     def test_check_grad(self):
         place = core.CUDAPlace(0)
@@ -280,6 +284,7 @@ def test_check_grad(self):
             'Out',
             user_defined_grads=[self.x_grad, self.y_grad],
             check_pir=True,
+            check_prim_pir=True,
         )
 
 
diff --git a/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py
@@ -110,6 +110,10 @@ def mean_all_net1(x):
     return paddle._C_ops.mean_all(x)
 
 
+def lerp_net(x, y, weight):
+    return paddle.lerp(x, y, weight)
+
+
 group_norm1 = paddle.nn.GroupNorm(num_channels=128, num_groups=32)
 
 
@@ -649,6 +653,127 @@ def setUp(self):
         self.enable_cinn = False
 
 
+class TestPrimThree(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [300, 2048]
+        self.shape_y = [300, 2048]
+        self.shape_z = [1]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.dtype_z = "float32"
+        self.init_x_shape = [None, 2048]
+        self.init_y_shape = [None, 2048]
+        self.init_z_shape = [None]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.y = np.random.random(self.shape_y).astype(self.dtype_y)
+        self.z = np.random.random(self.shape_z).astype(self.dtype_z)
+        self.net = lerp_net
+        self.necessary_ops = "pd_op.lerp"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+    def base_net(self, flag=None):
+        x = paddle.to_tensor(self.x)
+        y = paddle.to_tensor(self.y)
+        z = paddle.to_tensor(self.z)
+        if flag == "prim":
+            core._set_prim_all_enabled(True)
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=self.init_x_shape, dtype=self.dtype_x),
+                    InputSpec(shape=self.init_y_shape, dtype=self.dtype_y),
+                    InputSpec(shape=self.init_z_shape, dtype=self.dtype_z),
+                ],
+            )
+            fn.eval()
+        else:
+            fn = self.net
+        res = fn(x, y, z)
+
+        if flag == "prim":
+            ops = [
+                op.name()
+                for op in fn.program_cache.last()[-1][-1]
+                .infer_program.program.global_block()
+                .ops
+            ]
+            assert self.necessary_ops not in ops
+            core._set_prim_all_enabled(False)
+        return res
+
+    def test_prim_all_dynamic(self):
+        res_ref = self.base_net()
+        res = self.base_net("prim")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=self.tol)
+
+
+class TestPrimLerp1(TestPrimThree):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [10, 1, 10, 5, 5]
+        self.shape_y = [10, 5, 1, 5, 5]
+        self.shape_z = [1]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.dtype_z = "float32"
+        self.init_x_shape = [None, None, None, 5, 5]
+        self.init_y_shape = [None, None, None, 5, 5]
+        self.init_z_shape = [None]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.y = np.random.random(self.shape_y).astype(self.dtype_y)
+        self.z = np.random.random(self.shape_z).astype(self.dtype_z)
+        self.net = lerp_net
+        self.necessary_ops = "pd_op.lerp"
+        self.enable_cinn = False
+        self.tol = 1e-5
+
+
+class TestPrimLerp2(TestPrimThree):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [10, 10, 5, 5]
+        self.shape_y = [10, 10, 5, 5]
+        self.shape_z = [5]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.dtype_z = "float32"
+        self.init_x_shape = [None, None, 5, 5]
+        self.init_y_shape = [None, None, 5, 5]
+        self.init_z_shape = [None]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.y = np.random.random(self.shape_y).astype(self.dtype_y)
+        self.z = np.random.random(self.shape_z).astype(self.dtype_z)
+        self.net = lerp_net
+        self.necessary_ops = "pd_op.lerp"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimLerp3(TestPrimThree):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [10, 5, 10, 1, 5]
+        self.shape_y = [10, 5, 10, 5, 1]
+        self.shape_z = [1]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.dtype_z = "float32"
+        self.init_x_shape = [None, None, None, 1, 5]
+        self.init_y_shape = [None, None, None, 5, 1]
+        self.init_z_shape = [None]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.y = np.random.random(self.shape_y).astype(self.dtype_y)
+        self.z = np.random.random(self.shape_z).astype(self.dtype_z)
+        self.net = lerp_net
+        self.necessary_ops = "pd_op.lerp"
+        self.enable_cinn = False
+        self.tol = 1e-5
+
+
 class TestPrimLogLoss1(TestPrimTwo):
     def setUp(self):
         np.random.seed(2023)
