From 8c14b02b2f2281c0fe3d3d5fdb8d4972db45fb33 Mon Sep 17 00:00:00 2001 From: cyber-pioneer <116002591+cyber-pioneer@users.noreply.github.com> Date: Wed, 8 Feb 2023 15:07:10 +0800 Subject: [PATCH] [feature] use prim flag in shell (#50309) * add flag * change flag * use prim flag * fix code * fix softmax prim flag * set case timeout --- .../test_composite_softmax_grad.py | 8 +- .../unittests/prim/prim/flags/CMakeLists.txt | 5 + .../prim/prim/flags/test_prim_flags_case.py | 173 ++++++++++++++++++ .../paddle/jit/dy2static/partial_program.py | 5 +- .../jit/dy2static/program_translator.py | 5 +- 5 files changed, 187 insertions(+), 9 deletions(-) create mode 100644 python/paddle/fluid/tests/unittests/prim/prim/flags/test_prim_flags_case.py diff --git a/python/paddle/fluid/tests/unittests/prim/composite_ops/test_composite_softmax_grad.py b/python/paddle/fluid/tests/unittests/prim/composite_ops/test_composite_softmax_grad.py index 9b6e5db795..d8b3731332 100644 --- a/python/paddle/fluid/tests/unittests/prim/composite_ops/test_composite_softmax_grad.py +++ b/python/paddle/fluid/tests/unittests/prim/composite_ops/test_composite_softmax_grad.py @@ -78,7 +78,7 @@ class TestCompositeSoftmax(unittest.TestCase): def cal_composite_grad(self, inputs): paddle.enable_static() - core._set_prim_all_enabled(True) + core._set_prim_forward_enabled(True) startup_program = paddle.static.Program() main_program = paddle.static.Program() with paddle.static.program_guard(main_program, startup_program): @@ -109,7 +109,7 @@ class TestCompositeSoftmax(unittest.TestCase): exe.run(startup_program) res = exe.run(main_program, feed={'x': inputs}, fetch_list=[z]) paddle.disable_static() - core._set_prim_all_enabled(False) + core._set_prim_forward_enabled(False) return res def compare_backward(self): @@ -142,12 +142,13 @@ class TestCompositeSoftmaxPrimBackward(unittest.TestCase): def setUp(self): core._set_prim_backward_enabled(True) - self.dtypes = ["float32"] + self.dtypes = ["float32", "float64"] self.shapes = [[2, 3, 4], [2, 3]] self.axes = [-1, 0, 1] def cal_composite_grad(self, inputs): paddle.enable_static() + core._set_prim_all_enabled(True) startup_program = paddle.static.Program() main_program = paddle.static.Program() with paddle.static.program_guard(main_program, startup_program): @@ -164,6 +165,7 @@ class TestCompositeSoftmaxPrimBackward(unittest.TestCase): exe.run(startup_program) res = exe.run(main_program, feed={'x': inputs}, fetch_list=[z]) paddle.disable_static() + core._set_prim_all_enabled(False) return res def compare_backward(self): diff --git a/python/paddle/fluid/tests/unittests/prim/prim/flags/CMakeLists.txt b/python/paddle/fluid/tests/unittests/prim/prim/flags/CMakeLists.txt index 72c6bbd7d0..e57c6138d2 100644 --- a/python/paddle/fluid/tests/unittests/prim/prim/flags/CMakeLists.txt +++ b/python/paddle/fluid/tests/unittests/prim/prim/flags/CMakeLists.txt @@ -7,3 +7,8 @@ string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}") foreach(TEST_OP ${TEST_OPS}) py_test_modules(${TEST_OP} MODULES ${TEST_OP} ENVS ${GC_ENVS}) endforeach() + +if(WITH_CINN) + set_tests_properties(test_prim_flags_case PROPERTIES LABELS "RUN_TYPE=CINN") + set_tests_properties(test_prim_flags_case PROPERTIES TIMEOUT 300) +endif() diff --git a/python/paddle/fluid/tests/unittests/prim/prim/flags/test_prim_flags_case.py b/python/paddle/fluid/tests/unittests/prim/prim/flags/test_prim_flags_case.py new file mode 100644 index 0000000000..309959747e --- /dev/null +++ b/python/paddle/fluid/tests/unittests/prim/prim/flags/test_prim_flags_case.py @@ -0,0 +1,173 @@ +# Copyright (c) 2023 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 os +import platform +import unittest + +import paddle +import paddle.nn.functional as F +from paddle.fluid import core + + +def apply_to_static(net, use_cinn): + build_strategy = paddle.static.BuildStrategy() + build_strategy.build_cinn_pass = use_cinn + return paddle.jit.to_static(net, build_strategy=build_strategy) + + +class PrimeNet(paddle.nn.Layer): + def __init__(self): + super(PrimeNet, self).__init__() + + def forward(self, x): + out = F.softmax(x) + res = paddle.exp(out) + return res + + +class TestPrimForwardAndBackward(unittest.TestCase): + """ + Test PrimeNet with @to_static + prim forward + prim backward + cinn v.s Dygraph + """ + + def setUp(self): + paddle.seed(2022) + self.x = paddle.randn([2, 4]) + self.x.stop_gradient = False + self.flag = None + + def reset_env_flag(self): + os.environ["FLAGS_prim_backward"] = "False" + os.environ["FLAGS_prim_forward"] = "False" + if os.getenv("FLAGS_prim_all"): + del os.environ["FLAGS_prim_all"] + core.check_and_set_prim_all_enabled() + + def train(self, use_cinn): + net = PrimeNet() + net = apply_to_static(net, use_cinn) + + out = net(self.x) + loss = paddle.mean(out) + loss.backward() + + self.check_prim(net) + + return + + def check_prim(self, net): + ops = [ + op.type + for op in net.forward.program_cache.last()[-1][-1] + .train_program.block(0) + .ops + ] + + if self.flag in ["prim_all", "cinn_prim_all"]: + self.assertTrue('softmax' not in ops) + self.assertTrue('exp_grad' not in ops) + elif self.flag in ["prim_forward", "cinn_prim_forward"]: + self.assertTrue('softmax' not in ops) + self.assertTrue('exp_grad' in ops) + elif self.flag in ["prim_backward", "cinn_prim_backward"]: + self.assertTrue('softmax' in ops) + self.assertTrue('exp_grad' not in ops) + elif self.flag == "cinn": + self.assertTrue('softmax' in ops) + self.assertTrue('exp_grad' in ops) + else: + raise TypeError + + def test_cinn_prim_all(self): + """cinn + prim forward + prim backward""" + self.reset_env_flag() + os.environ["FLAGS_prim_all"] = "True" + self.flag = "cinn_prim_all" + plat = platform.system() + if plat == "Linux": + _ = self.train(use_cinn=True) + else: + pass + + def test_prim_all(self): + """prim forward + prim backward""" + self.reset_env_flag() + os.environ["FLAGS_prim_all"] = "True" + self.flag = "prim_all" + plat = platform.system() + if plat == "Linux": + _ = self.train(use_cinn=False) + else: + pass + + def test_cinn_prim_forward(self): + """cinn + prim forward""" + + self.reset_env_flag() + + os.environ["FLAGS_prim_forward"] = "True" + self.flag = "cinn_prim_forward" + plat = platform.system() + if plat == "Linux": + _ = self.train(use_cinn=True) + else: + pass + + def test_prim_forward(self): + """only prim forward""" + self.reset_env_flag() + os.environ["FLAGS_prim_forward"] = "True" + self.flag = "prim_forward" + plat = platform.system() + if plat == "Linux": + _ = self.train(use_cinn=False) + else: + pass + + def test_cinn_prim_backward(self): + """cinn + prim_backward""" + self.reset_env_flag() + os.environ["FLAGS_prim_backward"] = "True" + self.flag = "cinn_prim_backward" + plat = platform.system() + if plat == "Linux": + _ = self.train(use_cinn=True) + else: + pass + + def test_prim_backward(self): + """only prim backward""" + self.reset_env_flag() + os.environ["FLAGS_prim_backward"] = "True" + self.flag = "prim_backward" + plat = platform.system() + if plat == "Linux": + _ = self.train(use_cinn=False) + else: + pass + + def test_cinn(self): + """only cinn""" + self.reset_env_flag() + self.flag = "cinn" + plat = platform.system() + if plat == "Linux": + _ = self.train(use_cinn=True) + else: + pass + + +if __name__ == '__main__': + unittest.main() diff --git a/python/paddle/jit/dy2static/partial_program.py b/python/paddle/jit/dy2static/partial_program.py index fd509d74e5..626bdab2f5 100644 --- a/python/paddle/jit/dy2static/partial_program.py +++ b/python/paddle/jit/dy2static/partial_program.py @@ -571,9 +571,8 @@ class PartialProgramLayer: targets.append(program.global_block().var(out.name)) if targets: - if self._build_strategy.build_cinn_pass: - # TODO(Jiabin): Change this to True if we need this to be default option - core.check_and_set_prim_all_enabled() + # TODO(CZ): later when use cinn, set_prim_all_enabled and check_and_set_prim_all_enabled will be set at else branch. + core.check_and_set_prim_all_enabled() backward.gradients(targets=targets, inputs=[]) start_idx = len(main_program.block(0).ops) + len(self._outputs.tolist()) diff --git a/python/paddle/jit/dy2static/program_translator.py b/python/paddle/jit/dy2static/program_translator.py index f32168858b..37b92bbd4d 100644 --- a/python/paddle/jit/dy2static/program_translator.py +++ b/python/paddle/jit/dy2static/program_translator.py @@ -1140,9 +1140,8 @@ class ProgramCache: enable_prim = cache_key.kwargs['build_strategy'].build_cinn_pass # NOTE(xiongkun): Need a global FLAGS to enable/disable fallback enable_fallback = enable_prim - if enable_prim: - # TODO(Jiabin): Change this to True if we need this to be default option - core.check_and_set_prim_all_enabled() + # TODO(CZ): later when use cinn, set_prim_all_enabled and check_and_set_prim_all_enabled will be set at else branch. + core.check_and_set_prim_all_enabled() try: concrete_program = ConcreteProgram.from_func_spec( func_spec=cache_key.function_spec, -- GitLab