Skip to content

Commit c0b2c7d

Browse files
[Prim][PIR] support floor_divide op forward in prim pir (#64023)
* floor-div-dev * update test
1 parent a21ef65 commit c0b2c7d

File tree

3 files changed

+13
-1
lines changed

3 files changed

+13
-1
lines changed

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
"elu",
3030
"embedding",
3131
"flatten",
32+
"floor_divide",
3233
"full_like",
3334
"gelu",
3435
"hardswish",
@@ -67,6 +68,7 @@
6768
"elu",
6869
"embedding",
6970
"flatten",
71+
"floor_divide",
7072
"full_like",
7173
"gelu",
7274
"hardswish",

paddle/fluid/primitive/composite/composite.h

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -729,6 +729,14 @@ Tensor full_like_decomp(const Tensor& x,
729729
}
730730
}
731731

732+
template <typename T>
733+
Tensor floor_divide_decomp(const Tensor& x, const Tensor& y) {
734+
auto x_cast = cast<T>(x, DataType::INT64);
735+
auto y_cast = cast<T>(y, DataType::INT64);
736+
auto res = x_cast / y_cast;
737+
return cast<T>(res, x.dtype());
738+
}
739+
732740
template <typename T>
733741
std::tuple<Tensor, Tensor> dropout_decomp(
734742
const Tensor& x,

test/deprecated/legacy_test/test_elementwise_floordiv_op.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,9 @@ def init_kernel_type(self):
2929

3030
def setUp(self):
3131
self.op_type = "elementwise_floordiv"
32+
self.prim_op_type = "comp"
3233
self.python_api = paddle.floor_divide
34+
self.public_python_api = paddle.floor_divide
3335
self.dtype = np.int32
3436
self.axis = -1
3537
self.init_dtype()
@@ -45,7 +47,7 @@ def setUp(self):
4547
self.outputs = {'Out': self.out}
4648

4749
def test_check_output(self):
48-
self.check_output(check_pir=True)
50+
self.check_output(check_pir=True, check_prim_pir=True)
4951

5052
def init_input_output(self):
5153
self.x = np.random.uniform(0, 10000, [10, 10]).astype(self.dtype)

0 commit comments

Comments
 (0)