[dynamo] expose code execution strategy to python (pytorch#148020)

@anijain2305 this can be used to mark a code object to be skipped/run-only (recursively) while tracing.

Pull Request resolved: pytorch#148020
Approved by: https://github.com/jansel

Author: williamwen42

torch/_C/_dynamo/eval_frame.pyi

@@ -9,7 +9,9 @@ def set_skip_guard_eval_unsafe(value: bool) -> bool: ...
 def get_eval_frame_callback() -> DynamoCallback: ...
 def reset_code(code: types.CodeType) -> None: ...
 def unsupported(obj1: object, obj2: object) -> object: ...
-def skip_code(code: types.CodeType) -> None: ...
+def set_code_exec_strategy(
+    code: types.CodeType, strategy: _FrameExecStrategy
+) -> None: ...
 def set_guard_error_hook(hook: DynamoGuardHook) -> None: ...
 def raise_sigtrap() -> None: ...
 

torch/_dynamo/decorators.py

@@ -26,6 +26,7 @@
     DynamoStance,
     innermost_fn,
     RunOnlyContext,
+    skip_code,
 )
 from .exc import IncorrectUsage
 from .external_utils import is_compiling
@@ -39,7 +40,6 @@
         reset_code,
         set_eval_frame,
         set_guard_error_hook,
-        skip_code,
         unsupported,
     )
 

torch/_dynamo/eval_frame.py

@@ -56,10 +56,10 @@
 # see discussion at https://github.com/pytorch/pytorch/issues/120699
 from torch._C._dynamo.eval_frame import (  # noqa: F401
     reset_code,
+    set_code_exec_strategy,
     set_eval_frame,
     set_guard_error_hook,
     set_skip_guard_eval_unsafe,
-    skip_code,
     unsupported,
 )
 from torch._dispatch.python import enable_python_dispatcher
@@ -86,6 +86,7 @@
 from .exc import CondOpArgsMismatchError, ShortenTraceback, UserError, UserErrorType
 from .hooks import Hooks
 from .mutation_guard import install_generation_tagging_init
+from .types import FrameAction, FrameExecStrategy
 from .utils import common_constant_types, compile_times
 
 
@@ -1886,3 +1887,9 @@ def inner_fn(*args, **kwargs):
             return fn(*args, **kwargs)
 
         return inner_fn
+
+
+def skip_code(code: types.CodeType):
+    set_code_exec_strategy(
+        code, FrameExecStrategy(FrameAction.SKIP, FrameAction.DEFAULT)
+    )

torch/_dynamo/variables/functions.py

@@ -527,7 +527,7 @@ def next_variable(self, tx):
             # test/dynamo/test_misc.py::test_iterator_limit
             raise
         except Unsupported as e:
-            torch._C._dynamo.eval_frame.skip_code(self.get_code())
+            torch._dynamo.eval_frame.skip_code(self.get_code())
             raise SkipFrame from e
         finally:
             counters["unimplemented"] |= counters["inline_call"]

torch/csrc/dynamo/eval_frame.c

@@ -615,21 +615,6 @@ static PyObject* unsupported(PyObject* dummy, PyObject* args) {
   return obj2;
 }
 
-static PyObject* skip_code(PyObject* dummy, PyObject* obj) {
-  if (!PyCode_Check(obj)) {
-    PyErr_SetString(PyExc_TypeError, "expected a code object");
-    return NULL;
-  }
-
-  PyCodeObject* code = (PyCodeObject*)obj;
-  ExtraState* extra = get_extra_state(code);
-  if (extra == NULL) {
-    extra = init_and_set_extra_state(code);
-  }
-  extra_state_set_exec_strategy(extra, (FrameExecStrategy){SKIP, DEFAULT});
-  Py_RETURN_NONE;
-}
-
 static PyObject* set_guard_error_hook(PyObject* dummy, PyObject* obj) {
   if (obj == Py_None) {
     obj = NULL;
@@ -676,7 +661,7 @@ static PyMethodDef _methods[] = {
     {"get_eval_frame_callback", get_eval_frame_callback_py, METH_NOARGS, NULL},
     {"reset_code", reset_code, METH_O, NULL},
     {"unsupported", unsupported, METH_VARARGS, NULL},
-    {"skip_code", skip_code, METH_O, NULL},
+    {"set_code_exec_strategy", set_code_exec_strategy, METH_VARARGS, NULL},
     {"set_guard_error_hook", set_guard_error_hook, METH_O, NULL},
     {"raise_sigtrap", raise_sigtrap, METH_NOARGS, NULL},
     {NULL, NULL, 0, NULL}};

torch/csrc/dynamo/eval_frame_cpp.cpp

@@ -290,3 +290,27 @@ PyObject* dynamo__custom_eval_frame(
   }
   return eval_result;
 }
+
+PyObject* set_code_exec_strategy(PyObject* dummy, PyObject* args) {
+  PyObject* code_obj = nullptr;
+  PyObject* strategy_obj = nullptr;
+  if (!PyArg_ParseTuple(args, "OO", &code_obj, &strategy_obj)) {
+    return nullptr;
+  }
+  if (!PyCode_Check(code_obj)) {
+    PyErr_SetString(PyExc_TypeError, "expected a code object");
+    return nullptr;
+  }
+
+  PyCodeObject* code = (PyCodeObject*)code_obj;
+  ExtraState* extra = get_extra_state(code);
+  if (extra == nullptr) {
+    extra = init_and_set_extra_state(code);
+  }
+
+  FrameExecStrategy strategy =
+      py::handle(strategy_obj).cast<FrameExecStrategy>();
+
+  extra_state_set_exec_strategy(extra, strategy);
+  Py_RETURN_NONE;
+}

torch/csrc/dynamo/eval_frame_cpp.h

@@ -1,6 +1,5 @@
 #pragma once
 #include <Python.h>
-#include <stdbool.h>
 
 #include <torch/csrc/dynamo/eval_frame.h>
 #include <torch/csrc/dynamo/extra_state.h>
@@ -17,6 +16,8 @@ PyObject* dynamo__custom_eval_frame(
     int throw_flag,
     PyObject* callback);
 
+PyObject* set_code_exec_strategy(PyObject* dummy, PyObject* obj);
+
 #ifdef __cplusplus
 
 } // extern "C"