[Prim][PIR] fallback mechanism to support dropout op (#66711)

* add a fallback mechanism to support dropout op

* add fallback unittest

* Rename test file and update CMakeLists.txt

* remove unnecessary comments

* rename class method

Author: cubehan3

paddle/fluid/primitive/base/decomp_trans.cc

@@ -480,6 +480,10 @@ void DecompProgram::decomp_block(
       auto& builder = *(paddle::dialect::ApiBuilder::Instance().GetBuilder());
       builder.set_insertion_point(op);
       std::vector<std::vector<pir::Value>> decomp_res = call_decomp_rule(op);
+      if (decomp_res.size() == 0) {
+        // if we don't decomp this op, then leave it intact.
+        continue;
+      }
       std::vector<pir::Value> orig_outs = op->results();
       bool is_next_builtin_split_slice = false;
 

paddle/fluid/primitive/codegen/templates/decomp/generated_decomp.j2

@@ -81,10 +81,7 @@ std::vector<std::vector<pir::Value>> {{class_name}}::Decomp(pir::Operation* op)
           ->value()
           .defining_op();
   if ({{item.name}}_define_op->name() != "pd_op.full") {
-    PADDLE_THROW(
-        common::errors::Unimplemented("We don't support dynamic tensors "
-                                        "attribute {{item.name}} for {{fwd_name}} decomposition "
-                                        "for now. "));
+    return {};
   }
   Scalar {{item.name}} = {{item.name}}_define_op->attribute("value").dyn_cast<paddle::dialect::ScalarAttribute>().data();
 
@@ -97,10 +94,7 @@ std::vector<std::vector<pir::Value>> {{class_name}}::Decomp(pir::Operation* op)
           ->value()
           .defining_op();
   if ({{item.name}}_define_op->name() != "pd_op.full_int_array") {
-    PADDLE_THROW(
-        common::errors::Unimplemented("We don't support dynamic tensors "
-                                        "attribute {{item.name}} for {{fwd_name}} decomposition "
-                                        "for now. "));
+    return {};
   }
   IntArray {{item.name}} = phi::IntArray(
       paddle::dialect::GetInt64Vector({{item.name}}_define_op->attribute("value")));
@@ -120,7 +114,7 @@ std::vector<std::vector<pir::Value>> {{class_name}}::Decomp(pir::Operation* op)
       }
 
     } else {
-      PADDLE_THROW(common::errors::Unimplemented("attr is not vector of {{temp_type}} "));
+      return {};
     }
   }
         {% else %}

paddle/fluid/primitive/composite/composite.h

@@ -856,7 +856,7 @@ std::tuple<Tensor, Tensor> dropout_decomp(
       // train: out = input * mask / ( 1.0 - p )
       if (p.to<float>() == 1.0) {
         // Process p=1. for avoid divide zero error (x*mask/(1.0-p))
-        auto zero = full_scalar<T>(0.0, org_dtype);
+        auto zero = full_like_decomp<T>(x, 0.0, org_dtype, x.place());
         return std::make_tuple(x * zero, cast<T>(zero, DataType::UINT8));
       } else {
         auto ans = (x * mask) / ones_p;

test/prim/pir_prim/CMakeLists.txt

@@ -22,7 +22,8 @@ set(TEST_PRIM_PURE_PIR_CASES
     test_decompose_control_flow
     test_decomp_whole_program
     test_dynamic_combine1
-    test_dynamic_combine2)
+    test_dynamic_combine2
+    test_decomp_fallback)
 
 foreach(target ${TEST_PRIM_PURE_PIR_CASES})
   py_test_modules(

test/prim/pir_prim/test_decomp_fallback.py

@@ -0,0 +1,152 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+
+
+class TestFallBackBase(unittest.TestCase):
+    def setUp(self):
+        self.func_api = None
+        self.dtype = np.float32
+        self.tol = 1e-6
+
+
+def custom_dropout(x, p):
+    return paddle.nn.functional.dropout(x, p) + 2.0
+
+
+class TestDropOutFallBack(TestFallBackBase):
+    def setUp(self):
+        super().setUp()
+        self.func_api = custom_dropout
+        self.x = paddle.to_tensor([[1.0, -2], [3.0, 4]], dtype=self.dtype)
+        self.p = paddle.to_tensor(0.0, dtype=self.dtype)
+
+    def test_fallback(self):
+        static_func = paddle.jit.to_static(self.func_api, full_graph=True)
+        dynamic_func = self.func_api
+
+        out = static_func(self.x, self.p)
+        ref_out = dynamic_func(self.x, self.p)
+
+        for ref, actual in zip(ref_out, out):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+
+def custom_full(shape, value):
+    return paddle.full_like(shape, value) + 2.0
+
+
+class TestFullLikeFallBack(TestFallBackBase):
+    def setUp(self):
+        super().setUp()
+        self.func_api = custom_full
+        self.x = paddle.to_tensor([[1.0, -2], [3.0, 4]], dtype=self.dtype)
+        self.value = paddle.to_tensor(2, dtype=self.dtype)
+
+    def test_fallback(self):
+        static_func = paddle.jit.to_static(self.func_api, full_graph=True)
+        dynamic_func = self.func_api
+
+        out = static_func(self.x, self.value)
+        ref_out = dynamic_func(self.x, self.value)
+
+        for ref, actual in zip(ref_out, out):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+
+def custom_squeeze(x, axis):
+    return paddle.squeeze(x, axis) + 2.0
+
+
+class TestSqueezeFallBack(TestFallBackBase):
+    def setUp(self):
+        super().setUp()
+        self.func_api = custom_squeeze
+        self.x = paddle.rand([5, 1, 10], dtype=self.dtype)
+        self.axis = paddle.to_tensor(1, dtype=paddle.int64)
+
+    def test_fallback(self):
+        static_func = paddle.jit.to_static(self.func_api, full_graph=True)
+        dynamic_func = self.func_api
+
+        out = static_func(self.x, self.axis)
+        ref_out = dynamic_func(self.x, self.axis)
+
+        for ref, actual in zip(ref_out, out):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+
+def custom_unsqueeze(x, axis):
+    return paddle.unsqueeze(x, axis) + 2.0
+
+
+class TestUnsqueezeFallBack(TestFallBackBase):
+    def setUp(self):
+        super().setUp()
+        self.func_api = custom_unsqueeze
+        self.x = paddle.rand([5, 10], dtype=self.dtype)
+        self.axis = paddle.to_tensor([0, 2], dtype=paddle.int64)
+
+    def test_fallback(self):
+        static_func = paddle.jit.to_static(self.func_api, full_graph=True)
+        dynamic_func = self.func_api
+
+        out = static_func(self.x, self.axis)
+        ref_out = dynamic_func(self.x, self.axis)
+
+        for ref, actual in zip(ref_out, out):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+
+def custom_any(x, axis):
+    return paddle.any(x, axis)
+
+
+class TestAnyFallBack(TestFallBackBase):
+    def setUp(self):
+        super().setUp()
+        self.func_api = custom_any
+        self.x = paddle.to_tensor([[1, 0], [1, 1]], dtype='int32').cast('bool')
+        # Axis cannot accept a list of tensors,
+        # the framework will check the argument type before decomposition.
+        self.axis = [0]
+
+    def test_fallback(self):
+        static_func = paddle.jit.to_static(self.func_api, full_graph=True)
+        dynamic_func = self.func_api
+
+        out = static_func(self.x, self.axis)
+        ref_out = dynamic_func(self.x, self.axis)
+
+        for ref, actual in zip(ref_out, out):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()