wehub-resource-sync
119 lines
4.1 KiB
Python
119 lines
4.1 KiB
Python
# 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.
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, TypeVar
|
|
|
|
from typing_extensions import ParamSpec
|
|
|
|
import paddle
|
|
|
|
from .opcode_translator import eval_frame_callback
|
|
from .profiler import SotStepProfilerGuard
|
|
from .utils import (
|
|
InfoCollector,
|
|
StepInfoManager,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
P = ParamSpec("P")
|
|
R = TypeVar("R")
|
|
|
|
|
|
def symbolic_translate(fn: Callable[P, R], **kwargs) -> Callable[P, R]:
|
|
"""
|
|
This function is the entry point of PaddleSOT. It sets eval_frame_callback before input
|
|
function to achieve Opcode-level translation. The translation process depends on the
|
|
simulation execution, in which information will be collected, especially the network
|
|
code. After the simulation execution is completed, the network code will be compiled
|
|
into a static graph Program to improve performance.
|
|
|
|
Args:
|
|
fn: The input function.
|
|
|
|
Returns:
|
|
Callable, The wrapped function.
|
|
|
|
Examples:
|
|
>>> # doctest: +SKIP("Could not get source code of function foo."")
|
|
>>> import paddle
|
|
>>> import numpy as np
|
|
>>> from sot.translate import symbolic_translate
|
|
>>> def foo(cond: paddle.Tensor, x: paddle.Tensor):
|
|
... x += 1
|
|
... if cond:
|
|
... x += 1
|
|
... else:
|
|
... x -= 1
|
|
... return x
|
|
>>> symbolic_translate_foo = symbolic_translate(foo)
|
|
>>> # For the true branch, the output is 2.
|
|
>>> cond = paddle.to_tensor(True)
|
|
>>> x = paddle.to_tensor(0)
|
|
>>> dygraph_out = foo(cond, x)
|
|
>>> symbolic_translate_out = symbolic_translate_foo(cond, x)
|
|
>>> dygraph_out
|
|
Tensor(shape=[], dtype=int64, place=Place(cpu), stop_gradient=True,
|
|
2)
|
|
>>> symbolic_translate_out
|
|
Tensor(shape=[], dtype=int64, place=Place(cpu), stop_gradient=True,
|
|
2)
|
|
>>> np.testing.assert_allclose(dygraph_out.numpy(), symbolic_translate_out.numpy())
|
|
>>> # For the false branch, the output is 0.
|
|
>>> cond = paddle.to_tensor(False)
|
|
>>> dygraph_out = foo(cond, x)
|
|
>>> symbolic_translate_out = symbolic_translate_foo(cond, x)
|
|
>>> dygraph_out
|
|
Tensor(shape=[], dtype=int64, place=Place(cpu), stop_gradient=True,
|
|
0)
|
|
>>> symbolic_translate_out
|
|
Tensor(shape=[], dtype=int64, place=Place(cpu), stop_gradient=True,
|
|
0)
|
|
>>> np.testing.assert_allclose(dygraph_out.numpy(), symbolic_translate_out.numpy())
|
|
|
|
"""
|
|
|
|
if not paddle.framework.use_pir_api():
|
|
raise RuntimeError(
|
|
"SOT is only supported when running in PIR mode. Please set the environment variable "
|
|
"FLAGS_enable_pir_api=1 to enable it."
|
|
)
|
|
|
|
kwargs.setdefault('training', True)
|
|
|
|
def callback(frame):
|
|
return eval_frame_callback(frame, **kwargs)
|
|
|
|
def impl(*args: P.args, **kwargs: P.kwargs) -> R:
|
|
assert hasattr(fn, "__code__"), (
|
|
"Target function doesn't have code for simulating."
|
|
)
|
|
with StepInfoManager().step_guard(fn.__code__), SotStepProfilerGuard():
|
|
InfoCollector().clear_step_info()
|
|
paddle.framework.core.set_eval_frame(callback)
|
|
try:
|
|
outs = fn(*args, **kwargs)
|
|
except Exception as e:
|
|
raise e
|
|
finally:
|
|
paddle.framework.core.set_eval_frame(None)
|
|
|
|
InfoCollector().print_step_report()
|
|
return outs
|
|
|
|
return impl
|