[PIR] No.43 Migrate paddle.where into pir (PaddlePaddle#57667)

jiahy0825 · Oct 16, 2023 · 305e2d8 · 305e2d8
1 parent e6ac3a0
commit 305e2d8
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 4 deletions.
diff --git a/paddle/fluid/pir/dialect/op_generator/vjp_interface_gen_op_list.py b/paddle/fluid/pir/dialect/op_generator/vjp_interface_gen_op_list.py
@@ -24,6 +24,7 @@
 
 
 vjp_interface_declare_gen_op_list = [
+    'where',
     "tanh",
     "mean",
     "divide",
@@ -66,6 +67,7 @@
     'triu',
 ]
 vjp_interface_implementation_gen_op_list = [
+    'where',
     "tanh",
     "mean",
     "divide",

diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -38,6 +38,7 @@
 
 
 VJPS = [
+    'where_grad',
     'tril_grad',
     'triu_grad',
     'tanh_grad',
@@ -152,6 +153,7 @@
 VJP_COMPS = PRIM_VJP + CUSTOM_VJP
 
 BACKENDS = [
+    'where_grad',
     'tril_grad',
     'triu_grad',
     'add_n',

diff --git a/python/paddle/tensor/search.py b/python/paddle/tensor/search.py
@@ -27,6 +27,7 @@
     convert_np_dtype_to_dtype_,
     core,
     in_dynamic_mode,
+    in_dynamic_or_pir_mode,
 )
 
 # from ..base.layers import has_inf  #DEFINE_ALIAS
@@ -686,7 +687,7 @@ def where(condition, x=None, y=None, name=None):
         broadcast_condition = paddle.add(cast_cond, broadcast_zeros)
         broadcast_condition = paddle.cast(broadcast_condition, 'bool')
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.where(broadcast_condition, broadcast_x, broadcast_y)
     else:
         check_variable_and_dtype(condition, 'condition', ['bool'], 'where')

diff --git a/test/legacy_test/test_where_op.py b/test/legacy_test/test_where_op.py
@@ -33,10 +33,12 @@ def setUp(self):
         self.outputs = {'Out': np.where(self.cond, self.x, self.y)}
 
     def test_check_output(self):
-        self.check_output(check_cinn=self.check_cinn)
+        self.check_output(check_cinn=self.check_cinn, check_new_ir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X', 'Y'], 'Out', check_cinn=self.check_cinn)
+        self.check_grad(
+            ['X', 'Y'], 'Out', check_cinn=self.check_cinn, check_new_ir=True
+        )
 
     def init_config(self):
         self.x = np.random.uniform((-3), 5, 100).astype('float64')
@@ -82,7 +84,9 @@ def setUp(self):
 
     def test_check_output(self):
         place = core.CUDAPlace(0)
-        self.check_output_with_place(place, check_cinn=self.check_cinn)
+        self.check_output_with_place(
+            place, check_cinn=self.check_cinn, check_new_ir=True
+        )
 
     def test_check_grad(self):
         place = core.CUDAPlace(0)
@@ -92,6 +96,7 @@ def test_check_grad(self):
             'Out',
             numeric_grad_delta=0.05,
             check_cinn=self.check_cinn,
+            check_new_ir=True,
         )
 
     def init_config(self):