PaddlePaddle
diff --git a/‎python/paddle/base/dygraph/base.py
Copy file name to clipboard
+12-4Lines changed: 12 additions & 4 deletions b/‎python/paddle/base/dygraph/base.py
Copy file name to clipboard
+12-4Lines changed: 12 additions & 4 deletions
diff --git a/‎python/paddle/jit/sot/opcode_translator/executor/function_graph.py
Copy file name to clipboard
+32-18Lines changed: 32 additions & 18 deletions b/‎python/paddle/jit/sot/opcode_translator/executor/function_graph.py
Copy file name to clipboard
+32-18Lines changed: 32 additions & 18 deletions
diff --git a/‎python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
Copy file name to clipboard
+29Lines changed: 29 additions & 0 deletions b/‎python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
Copy file name to clipboard
+29Lines changed: 29 additions & 0 deletions
@@ -94,7 +94,7 @@ def to_static_unsupported_argument_warning(
 
 
 def _switch_to_static_graph_(
-    func: Callable[_InputT, _RetT]
+    func: Callable[_InputT, _RetT],
 ) -> Callable[_InputT, _RetT]:
     def __impl__(*args: _InputT.args, **kwargs: _InputT.kwargs) -> _RetT:
         with framework._dygraph_guard(None):
@@ -134,7 +134,7 @@ def sot_simulation_mode_guard(
 
 @signature_safe_contextmanager
 def param_guard(
-    parameters: OrderedDict[str, Tensor]
+    parameters: OrderedDict[str, Tensor],
 ) -> Generator[None, None, None]:
     # Note: parameters is a reference of self._parameters or self._buffers
     if in_to_static_mode() and not paddle.in_dynamic_mode() and parameters:
@@ -391,6 +391,8 @@ def __impl__(
 class _DecoratorContextManager:
     """Allow a context manager to be used as a decorator"""
 
+    DECORATED_BY_MARKER_ATTR = "__decorated_by__"
+
     def __call__(
         self, func: Callable[_InputT, _RetT]
     ) -> Callable[_InputT, _RetT]:
@@ -406,9 +408,15 @@ def _decorate_generator(func, *args, **kwargs):
                 yield from gen
 
         if inspect.isgeneratorfunction(func):
-            return _decorate_generator(func)
+            decorated_fn = _decorate_generator(func)
         else:
-            return _decorate_function(func)
+            decorated_fn = _decorate_function(func)
+        setattr(
+            decorated_fn,
+            _DecoratorContextManager.DECORATED_BY_MARKER_ATTR,
+            self,
+        )
+        return decorated_fn
 
     def __enter__(self) -> Any:
         raise NotImplementedError
 
@@ -19,6 +19,7 @@
 
 import inspect
 from collections import namedtuple
+from contextlib import contextmanager
 from copy import deepcopy
 from functools import reduce
 from typing import TYPE_CHECKING, Any, Callable, Tuple, Union
@@ -37,8 +38,14 @@
     ast_infer_meta,
 )
 from ...profiler import EventGuard, event_register
-from ...symbolic.statement_ir import Reference, StatementIR, Symbol
-from ...symbolic.symbolic_context import SymbolicTraceContext
+from ...symbolic.builder import StatementIRBuilder
+from ...symbolic.statement_ir import (
+    Reference,
+    StatementContext,
+    StatementContextRegistry,
+    StatementIR,
+    Symbol,
+)
 from ...symbolic_shape.operators import SYMBOLIC_BINARY_OPS, SYMBOLIC_UNARY_OPS
 from ...utils import (
     ENV_SOT_ALLOW_DYNAMIC_SHAPE,
@@ -199,6 +206,11 @@ def load(self, var):
             self._pycode_gen.gen_load(self._store_var_info[var.id][0])
 
 
+# TODO(SigureMo): Add AMP auto_guard context
+@StatementContextRegistry.register_context(lambda _: paddle.no_grad())
+class NoGradContext(StatementContext): ...
+
+
 class FunctionGraph:
     """
     A Graph representation corresponding to each FunctionFrame
@@ -224,7 +236,7 @@ class FunctionGraph:
     def __init__(
         self, code: types.CodeType, globals: dict[str, object], **kwargs
     ):
-        self.sir_ctx = SymbolicTraceContext()
+        self.sir_builder = StatementIRBuilder()
         self.inner_out = set()
         self.input_variables = []  # Store variables required within a function
         self.pycode_gen = PyCodeGen(code, globals, disable_eval_frame=True)
@@ -269,7 +281,7 @@ def save_memo(self) -> FunctionGraph.Memo:
         NOTE:
             Why don't use __deepcopy__, because memo is not a deepcopy, i.e inner_out is only a shallow copy, SIR is a deepcopy.
         """
-        saved_stmt_ir = deepcopy(self.sir_ctx.TOS)
+        saved_stmt_ir = deepcopy(self.sir_builder.current_sir)
         return FunctionGraph.Memo(
             inner_out=set(self.inner_out),
             input_variables=list(self.input_variables),
@@ -291,7 +303,7 @@ def restore_memo(self, memo: FunctionGraph.Memo):
         """
         self.inner_out = memo.inner_out
         self.input_variables = memo.input_variables
-        self.sir_ctx.replace_TOS(memo.stmt_ir)
+        self.sir_builder.replace_current_sir(memo.stmt_ir)
         self._global_guarded_variables = memo.global_guards
         self.side_effects.restore_state(memo.side_effects_state)
         self._print_variables = memo.print_variables
@@ -442,19 +454,19 @@ def compile_graph(self, *ret_vars: VariableBase) -> CompileGraphResult:
         ]
 
         symbolic_outputs = self._find_tensor_outputs(ret_items)
-        statement_ir = self.sir_ctx.return_TOS(
+        statement_ir = self.sir_builder.finalize(
             [Symbol(tensor_var.var_name) for tensor_var in symbolic_outputs]
         )
         if not statement_ir.statements:
-            return self.sir_ctx.compile_do_nothing(), (
+            return self.sir_builder.compile_do_nothing(), (
                 statement_ir,
                 OrderedSet(),
                 OrderedSet(),
             )
         SIRToCodeMap().register(statement_ir, self.pycode_gen._origin_code)
         input_names = statement_ir.inputs
         symbolic_inputs = self._find_tensor_inputs(input_names)
-        compiled_fn = self.sir_ctx.compile_fn(
+        compiled_fn = self.sir_builder.compile_fn(
             statement_ir.name,
             tuple(var.meta.to_input_spec() for var in symbolic_inputs),
             **self._kwargs,
@@ -520,17 +532,14 @@ def call_paddle_api(
             func: paddle api
         """
         assert is_paddle_api(func)
-        # not fallback api, start symbolic trace.
-        # TODO(xiokgun): may have python builtin object inside metas.
-        # TODO(xiokgun): 4 kinds of python arguments. support it !!
         log(3, f"call paddle.api : {func.__name__}", "\n")
 
         def message_handler(*args, **kwargs):
             return f"Call paddle_api error: {func.__name__}, may be not a operator api?"
 
         return inner_error_default_handler(self.symbolic_call, message_handler)(
             InferMetaCache(),
-            self.sir_ctx.call_API,
+            self.sir_builder.call_API,
             func,
             False,
             *args,
@@ -551,7 +560,7 @@ def message_handler(*args, **kwargs):
 
         return inner_error_default_handler(self.symbolic_call, message_handler)(
             InferMetaCache(),
-            self.sir_ctx.call_API,
+            self.sir_builder.call_API,
             op,
             True,
             *args,
@@ -573,7 +582,7 @@ def message_handler(*args, **kwargs):
 
         return inner_error_default_handler(self.symbolic_call, message_handler)(
             InferMetaCache(),
-            self.sir_ctx.call_METHOD,
+            self.sir_builder.call_METHOD,
             method_name,
             False,
             *args,
@@ -599,7 +608,7 @@ def infer_meta_fn(layer, *metas, **kwmetas):
             return metas
 
         def compute_fn(layer, inputs, outputs, stacks):
-            self.sir_ctx.call_LAYER(
+            self.sir_builder.call_LAYER(
                 Reference(layer.value, weak_ref),
                 inputs=inputs,
                 outputs=outputs,
@@ -627,7 +636,7 @@ def call_ast(
         """
 
         def compute_fn(static_function, inputs, outputs, stacks):
-            self.sir_ctx.call_AST(
+            self.sir_builder.call_AST(
                 static_function,
                 inputs=inputs,
                 outputs=outputs,
@@ -750,7 +759,7 @@ def try_infer_meta_fn(args, kwargs) -> Any:
             convert_to_symbol(kwargs),
         )
 
-        record_symbols(self.sir_ctx.TOS, *args, **kwargs)
+        record_symbols(self.sir_builder.current_sir, *args, **kwargs)
 
         log(3, f"         inputs : {inputs_symbols}", "\n")
 
@@ -819,13 +828,18 @@ def add_alias(
         alias_fn = lambda x: x
         alias_fn.__name__ = "__sir_alias__"
         inputs_arg_pack = ([src], {})
-        self.sir_ctx.call_API(
+        self.sir_builder.call_API(
             alias_fn,
             convert_to_symbol(inputs_arg_pack),
             convert_to_symbol(dst),
             [],
         )
 
+    @contextmanager
+    def no_grad(self):
+        with self.sir_builder.attach_statement_context_guard(NoGradContext()):
+            yield
+
     @staticmethod
     def get_opcode_executor_stack():
         # NOTE: only for debug.
 
@@ -441,7 +441,7 @@ def print_sir(self):
         Prints the Static Instruction Representation (SIR) in the executor.
 
         """
-        print(self._graph.sir_ctx.TOS)
+        print(self._graph.sir_builder.current_sir)
 
     def _prepare_virtual_env(self):
         """
 
@@ -29,6 +29,7 @@
 )
 
 import paddle
+from paddle.base.dygraph.base import _DecoratorContextManager
 
 from .... import psdb
 from ....profiler import EventGuard
@@ -912,6 +913,34 @@ def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
         return None
 
 
+class NoGradFunctionVariable(FunctionVariable):
+    def __init__(
+        self, fn: Callable[..., Any], graph: FunctionGraph, tracker: Tracker
+    ):
+        super().__init__(fn, graph, tracker)
+        self.value = fn
+
+    def call_function(self, /, *args, **kwargs):
+        wrapped_fn = self.value.__wrapped__
+        wrapped_fn = VariableFactory.from_value(
+            wrapped_fn, self.graph, GetAttrTracker(self, "__wrapped__")
+        )
+        with self.graph.no_grad():
+            return wrapped_fn(*args, **kwargs)
+
+    @VariableFactory.register_from_value(
+        successor="UserDefinedFunctionVariable"
+    )
+    def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
+        if (
+            decorated_by := getattr(
+                value, _DecoratorContextManager.DECORATED_BY_MARKER_ATTR, None
+            )
+        ) and (isinstance(decorated_by, paddle.no_grad)):
+            return NoGradFunctionVariable(value, graph, tracker)
+        return None
+
+
 class UserDefinedGeneratorFunctionVariable(FunctionVariable):
     """
     UserDefinedGeneratorFunctionVariable is a subclass of FunctionVariable used to wrap a user-defined generator.