[Prim][PIR] support floor_divide op forward in prim pir (#64023)

* floor-div-dev * update test
PaddlePaddle · May 7, 2024 · c0b2c7d · c0b2c7d
1 parent a21ef65
commit c0b2c7d
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 1 deletion.
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
@@ -29,6 +29,7 @@
     "elu",
     "embedding",
     "flatten",
+    "floor_divide",
     "full_like",
     "gelu",
     "hardswish",
@@ -67,6 +68,7 @@
     "elu",
     "embedding",
     "flatten",
+    "floor_divide",
     "full_like",
     "gelu",
     "hardswish",

diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -729,6 +729,14 @@ Tensor full_like_decomp(const Tensor& x,
   }
 }
 
+template <typename T>
+Tensor floor_divide_decomp(const Tensor& x, const Tensor& y) {
+  auto x_cast = cast<T>(x, DataType::INT64);
+  auto y_cast = cast<T>(y, DataType::INT64);
+  auto res = x_cast / y_cast;
+  return cast<T>(res, x.dtype());
+}
+
 template <typename T>
 std::tuple<Tensor, Tensor> dropout_decomp(
     const Tensor& x,

diff --git a/test/deprecated/legacy_test/test_elementwise_floordiv_op.py b/test/deprecated/legacy_test/test_elementwise_floordiv_op.py
@@ -29,7 +29,9 @@ def init_kernel_type(self):
 
     def setUp(self):
         self.op_type = "elementwise_floordiv"
+        self.prim_op_type = "comp"
         self.python_api = paddle.floor_divide
+        self.public_python_api = paddle.floor_divide
         self.dtype = np.int32
         self.axis = -1
         self.init_dtype()
@@ -45,7 +47,7 @@ def setUp(self):
         self.outputs = {'Out': self.out}
 
     def test_check_output(self):
-        self.check_output(check_pir=True)
+        self.check_output(check_pir=True, check_prim_pir=True)
 
     def init_input_output(self):
         self.x = np.random.uniform(0, 10000, [10, 10]).astype(self.dtype)