[Prim][Pir] Decomp addmm op (#68835)

eggman-1024 · web-flow · commit dda47a084fb8 · 2024-12-02T14:06:38.000+08:00
* Decomp addmm op

* add dynamic shape test file

* modified code

* record temp code

* modified addmm_decomp()

* modified test case
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
@@ -20,6 +20,7 @@
 # manual decomp interface declare are located in manual_op.h
 decomp_interface_declare_gen_op_list = [
     "add_n",
+    "addmm",
     "any",
     "batch_norm",
     "batch_norm_",
@@ -76,6 +77,7 @@
 decomp_interface_implementation_gen_op_list = [
     "any",
     "add_n",
+    "addmm",
     "bce_loss",
     "bmm",
     "dropout",
diff --git a/paddle/fluid/primitive/decomp_rule/decomp_rule/composite.h b/paddle/fluid/primitive/decomp_rule/decomp_rule/composite.h
@@ -1402,6 +1402,17 @@ Tensor swish_decomp(const Tensor& x) {
   return x * sigmoid<T>(x);
 }
 
+template <typename T>
+Tensor addmm_decomp(const Tensor& input,
+                    const Tensor& x,
+                    const Tensor& y,
+                    const float beta,
+                    const float alpha) {
+  Tensor x_y_mat = matmul<T>(x, y);
+  return full_scalar<T>(alpha, x_y_mat.dtype()) * x_y_mat +
+         full_scalar<T>(beta, input.dtype()) * input;
+}
+
 }  // namespace details
 
 }  // namespace primitive
diff --git a/test/legacy_test/test_addmm_op.py b/test/legacy_test/test_addmm_op.py
@@ -26,7 +26,9 @@ class TestAddMMOp(OpTest):
     # test basic
     def setUp(self):
         self.op_type = "addmm"
+        self.prim_op_type = "comp"
         self.python_api = paddle.addmm
+        self.public_python_api = paddle.addmm
         self.init_dtype_type()
         self.inputs = {
             'Input': np.random.random((100, 1)).astype(self.dtype),
@@ -42,19 +44,42 @@ def init_dtype_type(self):
         self.dtype = np.float64
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True, check_prim_pir=True)
 
     def test_check_grad_normal(self):
-        self.check_grad(['Input', 'X', 'Y'], 'Out')
+        self.check_grad(
+            ['Input', 'X', 'Y'],
+            'Out',
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
     def test_check_grad_x(self):
-        self.check_grad(['X'], 'Out', no_grad_set=None)
+        self.check_grad(
+            ['X'],
+            'Out',
+            no_grad_set=None,
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
     def test_check_grad_y(self):
-        self.check_grad(['Y'], 'Out', no_grad_set=None)
+        self.check_grad(
+            ['Y'],
+            'Out',
+            no_grad_set=None,
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
     def test_check_grad_input(self):
-        self.check_grad(['Input'], 'Out', no_grad_set=None)
+        self.check_grad(
+            ['Input'],
+            'Out',
+            no_grad_set=None,
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
 
 class TestAddMMFP16Op(TestAddMMOp):
@@ -219,7 +244,9 @@ class TestAddMMOp2(TestAddMMOp):
     # test alpha and beta
     def setUp(self):
         self.op_type = "addmm"
+        self.prim_op_type = "comp"
         self.python_api = paddle.addmm
+        self.public_python_api = paddle.addmm
         self.dtype = np.float64
         self.init_dtype_type()
         self.inputs = {
@@ -241,7 +268,9 @@ class TestAddMMOp3(OpTest):
     # test broadcast
     def setUp(self):
         self.op_type = "addmm"
+        self.prim_op_type = "comp"
         self.python_api = paddle.addmm
+        self.public_python_api = paddle.addmm
         self.dtype = np.float64
         self.init_dtype_type()
         self.inputs = {
@@ -262,26 +291,40 @@ def init_dtype_type(self):
         pass
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True, check_prim_pir=True)
 
     def test_check_grad_normal(self):
-        self.check_grad(['Input', 'X', 'Y'], 'Out')
+        self.check_grad(
+            ['Input', 'X', 'Y'], 'Out', check_pir=True, check_prim_pir=True
+        )
 
     def test_check_grad_x(self):
-        self.check_grad(['X'], 'Out', no_grad_set=None)
+        self.check_grad(
+            ['X'], 'Out', no_grad_set=None, check_pir=True, check_prim_pir=True
+        )
 
     def test_check_grad_y(self):
-        self.check_grad(['Y'], 'Out', no_grad_set=None)
+        self.check_grad(
+            ['Y'], 'Out', no_grad_set=None, check_pir=True, check_prim_pir=True
+        )
 
     def test_check_grad_input(self):
-        self.check_grad(['Input'], 'Out', no_grad_set=None)
+        self.check_grad(
+            ['Input'],
+            'Out',
+            no_grad_set=None,
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
 
 class TestAddMMOp4(OpTest):
     # test broadcast
     def setUp(self):
         self.op_type = "addmm"
+        self.prim_op_type = "comp"
         self.python_api = paddle.addmm
+        self.public_python_api = paddle.addmm
         self.dtype = np.float64
         self.init_dtype_type()
         self.inputs = {
@@ -302,19 +345,26 @@ def init_dtype_type(self):
         pass
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True, check_prim_pir=True)
 
     def test_check_grad_normal(self):
-        self.check_grad(['Input', 'X', 'Y'], 'Out')
+        self.check_grad(
+            ['Input', 'X', 'Y'], 'Out', check_pir=True, check_prim_pir=True
+        )
 
     def test_check_grad_x(self):
-        self.check_grad(['X'], 'Out', no_grad_set=None)
+        self.check_grad(['X'], 'Out', no_grad_set=None, check_pir=True)
 
     def test_check_grad_y(self):
-        self.check_grad(['Y'], 'Out', no_grad_set=None)
+        self.check_grad(['Y'], 'Out', no_grad_set=None, check_pir=True)
 
     def test_check_grad_input(self):
-        self.check_grad(['Input'], 'Out', no_grad_set=None)
+        self.check_grad(
+            ['Input'],
+            'Out',
+            no_grad_set=None,
+            check_pir=True,
+        )
 
 
 class TestAddMMOp5(unittest.TestCase):
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
@@ -22,6 +22,10 @@
 from paddle.static import InputSpec
 
 
+def addmm_net(input, x, y):
+    return paddle.addmm(input, x, y, alpha=1.0, beta=2.0)
+
+
 def apply_to_static(net, use_cinn, input_spec=None):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
@@ -912,6 +916,28 @@ def test_prim_all_dynamic(self):
             np.testing.assert_allclose(ref, actual, rtol=self.tol)
 
 
+class TestPrimAddmm(TestPrimThree):
+    def setUp(self):
+        np.random.seed(2024)
+        paddle.seed(2024)
+        self.shape_x = [30, 50]
+        self.shape_y = [30, 80]
+        self.shape_z = [80, 50]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.dtype_z = "float32"
+        self.init_x_shape = [None, None]
+        self.init_y_shape = [None, None]
+        self.init_z_shape = [None, 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 = addmm_net
+        self.necessary_ops = "pd_op.addmm"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimLerp1(TestPrimThree):
     def setUp(self):
         np.random.seed(2023)
