From 27aa648ba7abd911c62fe06bb3150d985ad799cc Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Sun, 7 Nov 2021 16:55:09 +0800 Subject: [PATCH] fix(mge): control compute_mode by context GitOrigin-RevId: f3968309f4f0db97b57fee1a65f54de7e366b5dc --- imperative/python/megengine/config.py | 3 + imperative/python/megengine/core/_config.py | 205 ++++++++++++++++++ .../megengine/functional/debug_param.py | 50 +++-- 3 files changed, 242 insertions(+), 16 deletions(-) create mode 100644 imperative/python/megengine/config.py create mode 100644 imperative/python/megengine/core/_config.py diff --git a/imperative/python/megengine/config.py b/imperative/python/megengine/config.py new file mode 100644 index 000000000..270bfd482 --- /dev/null +++ b/imperative/python/megengine/config.py @@ -0,0 +1,3 @@ +from .core._config import * + +__import__("mprop").init() diff --git a/imperative/python/megengine/core/_config.py b/imperative/python/megengine/core/_config.py new file mode 100644 index 000000000..70d663abe --- /dev/null +++ b/imperative/python/megengine/core/_config.py @@ -0,0 +1,205 @@ +# -*- coding: utf-8 -*- +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2021 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +import os +from contextlib import contextmanager + +__compute_mode = "default" +__conv_format = "default" +_benchmark_kernel = False +_deterministic_kernel = False +_async_level = os.getenv("MEGENGINE_INTERP_ASYNC_LEVEL", 2) + + +__all__ = [ + "benchmark_kernel", + "deterministic_kernel", + "async_level", + "_compute_mode", + "_conv_format", + "_override", +] + + +@property +def benchmark_kernel(mod): + r"""Whether or not run possible algorithms on real device to find the best one. The default option is false, + which means use heuristic to choose the fastest algorithm. + + Examples: + .. code-block:: + + import megengine as mge + mge.config.benchmark_kernel = True + """ + return _benchmark_kernel + + +@benchmark_kernel.setter +def benchmark_kernel(mod, option: bool): + global _benchmark_kernel + _benchmark_kernel = option + + +@property +def deterministic_kernel(mod): + r"""Whether or not the fastest algorithm choosed is reproducible. The default option is false, + which means the algorithm is not reproducible. + + Examples: + .. code-block:: + + import megengine as mge + mge.config.deterministic_kernel = True + """ + return _deterministic_kernel + + +@deterministic_kernel.setter +def deterministic_kernel(mod, option: bool): + global _deterministic_kernel + _deterministic_kernel = option + + +@property +def async_level(mod) -> int: + r"""Get or set config whether raise error exactly when invoking op. The default level is 2, + which means both device and user side errors are async. + + Examples: + .. code-block:: + + import megengine as mge + mge.config.async_level = 2 + """ + return _async_level + + +@async_level.setter +def async_level(mod, level: int): + global _async_level + _async_level = level + + +@property +def _compute_mode(mod): + r"""Get or set the precision of intermediate results. The default option is "default", + which means that no special requirements will be placed on. When set to 'float32', it + would be used for accumulator and intermediate result, but only effective when input and + output are of float16 dtype. + + Examples: + .. code-block:: + + import megengine as mge + mge.config._compute_mode = "default" + """ + return __compute_mode + + +@_compute_mode.setter +def _compute_mode(mod, _compute_mode: str): + global __compute_mode + __compute_mode = _compute_mode + + +@property +def _conv_format(mod): + r"""Get or set convolution data/filter/output layout format. The default option is "default", + which means that no special format will be placed on. There are all layout definitions + + ``NCHW`` layout: ``{N, C, H, W}`` + ``NHWC`` layout: ``{N, H, W, C}`` + ``NHWCD4`` layout: ``{N, H, (C + 3) / 4, W, 4}`` + ``NHWCD4I`` layout: with ``align_axis = 2`` + ``NCHW4`` layout: ``{N, C/4, H, W, 4}`` + ``NCHW88`` layout: ``{N, C/8, H, W, 8}`` + ``CHWN4`` layout: ``{C/4, H, W, N, 4}`` + ``NCHW64`` layout: ``{N, C/64, H, W, 64}`` + + Examples: + .. code-block:: + + import megengine as mge + mge.config._conv_format = "NHWC" + """ + return __conv_format + + +@_conv_format.setter +def _conv_format(mod, format: str): + global __conv_format + __conv_format = format + + +def _reset_execution_config( + benchmark_kernel=None, + deterministic_kernel=None, + async_level=None, + compute_mode=None, + conv_format=None, +): + global _benchmark_kernel, _deterministic_kernel, _async_level, __compute_mode, __conv_format + orig_flags = ( + _benchmark_kernel, + _deterministic_kernel, + _async_level, + __compute_mode, + __conv_format, + ) + if benchmark_kernel is not None: + _benchmark_kernel = benchmark_kernel + if deterministic_kernel is not None: + _deterministic_kernel = deterministic_kernel + if async_level is not None: + _async_level = async_level + if compute_mode is not None: + __compute_mode = compute_mode + if conv_format is not None: + __conv_format = conv_format + + return orig_flags + + +@contextmanager +def _override( + benchmark_kernel=None, + deterministic_kernel=None, + async_level=None, + compute_mode=None, + conv_format=None, +): + r"""A context manager that users can opt in by attaching the decorator to set + the config of the global variable. + + Examples: + .. code-block:: + + import megengine as mge + + @mge.config._override( + benchmark_kernel = True, + deterministic_kernel = Fasle, + async_level=2, + compute_mode="float32", + conv_format="NHWC", + ) + def train(): + """ + orig_flags = _reset_execution_config( + benchmark_kernel, deterministic_kernel, async_level, compute_mode, conv_format, + ) + try: + yield + finally: + # recover the previous values + _reset_execution_config(*orig_flags) + + +def _get_actual_op_param(function_param, config_param): + return function_param if config_param == "default" else config_param diff --git a/imperative/python/megengine/functional/debug_param.py b/imperative/python/megengine/functional/debug_param.py index 43d00596d..9cbf56909 100644 --- a/imperative/python/megengine/functional/debug_param.py +++ b/imperative/python/megengine/functional/debug_param.py @@ -8,26 +8,38 @@ # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. import os +from ..core import _config from ..core.ops import builtin from ..logger import get_logger from ..utils.deprecation import deprecated Strategy = builtin.ops.Convolution.Strategy -_execution_strategy = os.getenv("MEGENGINE_EXECUTION_STRATEGY", "HEURISTIC") - if os.getenv("MEGENGINE_CONV_EXECUTION_STRATEGY") != None: get_logger().warning( "Environment variable `MEGENGINE_CONV_EXECUTION_STRATEGY` is deprecated, please use `MEGENGINE_EXECUTION_STRATEGY`" ) +_valid_string_option = { + "REPRODUCIBLE": Strategy.REPRODUCIBLE, + "HEURISTIC": Strategy.HEURISTIC, + "PROFILE": Strategy.PROFILE, +} + def get_execution_strategy() -> Strategy: r"""Returns the execution strategy of :class:`~module..Conv2d` and :func:`~.matmul` See :func:`~.set_execution_strategy` for possible return values """ - return _execution_strategy + strategy = Strategy(0) + if _config._benchmark_kernel: + strategy |= Strategy.PROFILE + else: + strategy |= Strategy.HEURISTIC + if _config._deterministic_kernel: + strategy |= Strategy.REPRODUCIBLE + return strategy def set_execution_strategy(option): @@ -50,7 +62,6 @@ def set_execution_strategy(option): * 'HEURISTIC' uses heuristic to choose the fastest algorithm. * 'PROFILE' runs possible algorithms on real device to find the best one. - * 'PROFILE_HEURISTIC' uses profiling result and heuristic to choose the fastest algorithm. * 'PROFILE_REPRODUCIBLE' uses the fastest of profiling result that is also reproducible. * 'HEURISTIC_REPRODUCIBLE' uses heuristic to choose the fastest algorithm that is also reproducible. @@ -58,29 +69,33 @@ def set_execution_strategy(option): It can also be set through the environment variable 'MEGENGINE_EXECUTION_STRATEGY'. """ - valid_string_option = { - "REPRODUCIBLE": Strategy.REPRODUCIBLE, - "HEURISTIC": Strategy.HEURISTIC, - "PROFILE": Strategy.PROFILE, - } - global _execution_strategy # pylint: disable=global-statement if isinstance(option, Strategy): - _execution_strategy = option + _config._benchmark_kernel = ( + True if option & _valid_string_option["PROFILE"] != Strategy(0) else False + ) + _config._deterministic_kernel = ( + True + if option & _valid_string_option["REPRODUCIBLE"] != Strategy(0) + else False + ) return assert isinstance(option, str) - strategy_tmp = Strategy(0) + _config._benchmark_kernel = False + _config._deterministic_kernel = False for opt in option.split("_"): - if not opt in valid_string_option: + if not opt in _valid_string_option: raise ValueError( "Valid option can only be one of {}, or combine them with '_'.".format( - valid_string_option.keys() + _valid_string_option.keys() ) ) - strategy_tmp = strategy_tmp | valid_string_option[opt] - _execution_strategy = strategy_tmp + _config._benchmark_kernel |= _valid_string_option[opt] == Strategy.PROFILE + _config._deterministic_kernel |= ( + _valid_string_option[opt] == Strategy.REPRODUCIBLE + ) @deprecated(version="1.3", reason="use get_execution_strategy() instead") @@ -91,3 +106,6 @@ def get_conv_execution_strategy() -> str: @deprecated(version="1.3", reason="use set_execution_strategy() instead") def set_conv_execution_strategy(option: str): return set_execution_strategy(option) + + +set_execution_strategy(os.getenv("MEGENGINE_EXECUTION_STRATEGY", "HEURISTIC")) -- GitLab