CPython 톺아보기 - Tokenizer, AST, bytecode, PyCodeObject, PyFrameObject + TorchDynamo

TL;DR

메이저 스크립트 언어의 실체는 대부분 C++이며(Python, Javascript 등), PyTorch와 같은 Python deep learning framework을 사용하더라도 성능 최적화 시에는 결국 C++ 구현체를 사용하게 됩니다. Python은 C++로 짜여진 Stack Machine 이라 표현하기도 합니다.

우리가 입력한 Python 코드의 실행흐름을 주요 컴포넌트 관점에서 대략적으로 표현하면 아래와 같습니다.
: Python Code -> tokenizing -> AST -> python byte code(pyc) -> PyCodeObject -> PyFrameObject -> PyFrameEx

구문분석 이후 생성한 AST를 바탕으로 byte code를 생성하고 이를 PyCodeObject가 가지고 있게 됩니다. PyFrameEx는 PyCodeObject 의 byte code를 해석하여 instruction을 실행합니다.

과정에서 호출되는 주요 함수들에 대해선 아래 문서에 자세히 설명되어있습니다.
https://docs.python.org/3/c-api/veryhigh.html#c.PyEval_EvalFrameEx

The Very High Level Layer

이번 포스팅에서는 Python file 실행부터 결과 출력까지의 생성되는 주요 컴포넌트들을 코드 레벨에서 다뤄보고 이를 통해 전반적인 CPython 아키텍처를 파악해보고자 합니다.

Main

아래 명령어를 실행하게 되면

python main.py

컴파일된 c++ binary인 python을 실행하게 됩니다.

여기서 window의 경우 표준입력을 텍스트 모드로 처리하고, 그 외 운영체제는 바이트 모드로 처리하기 때문에 전처리기로 이용해 Py_Main, Py_BytesMain 두 함수를 구분하고 있습니다.

#include "Python.h"

#ifdef MS_WINDOWS
int
wmain(int argc, wchar_t **argv)
{
    return Py_Main(argc, argv);
}
#else
int
main(int argc, char **argv)
{
    return Py_BytesMain(argc, argv);
}
#endif

이후 config를 로드하고 실행모드에 따라 분기되어, 위 python main.py 와 같은 경우 pymain_run_file을 호출하게 됩니다. wrapper method는 다르지만 핵심 구현은 모두 동일합니다.

if (config->run_command) {
    *exitcode = pymain_run_command(config->run_command);
}
else if (config->run_module) {
    *exitcode = pymain_run_module(config->run_module, 1);
}
else if (main_importer_path != NULL) {
    *exitcode = pymain_run_module(L"__main__", 0);
}
else if (config->run_filename != NULL) {
    *exitcode = pymain_run_file(config);
}
else {
    *exitcode = pymain_run_stdin(config);
}

pymain_run_file을 조금 더 타고 들어가다보면 조금 흥미로운 부분을 확인할 수 있는데요, pyc file이 있으면 pyc file을 실행하고, 아니면 처음부터 실행하는 부분을 확인 할 수 있습니다.

int
_PyRun_SimpleFileObject(FILE *fp, PyObject *filename, int closeit,
                        PyCompilerFlags *flags)
{
    ~ 생략 ~ 

    int pyc = maybe_pyc_file(fp, filename, closeit);
    if (pyc < 0) {
        goto done;
    }

    PyObject *v;
    if (pyc) {
        FILE *pyc_fp;
        /* Try to run a pyc file. First, re-open in binary */
        if (closeit) {
            fclose(fp);
        }

        pyc_fp = _Py_fopen_obj(filename, "rb");
        if (pyc_fp == NULL) {
            fprintf(stderr, "python: Can't reopen .pyc file\n");
            goto done;
        }

        if (set_main_loader(dict, filename, "SourcelessFileLoader") < 0) {
            fprintf(stderr, "python: failed to set __main__.__loader__\n");
            ret = -1;
            fclose(pyc_fp);
            goto done;
        }
        v = run_pyc_file(pyc_fp, dict, dict, flags);
    } else {
        /* When running from stdin, leave __main__.__loader__ alone */
        if ((!PyUnicode_Check(filename) || !PyUnicode_EqualToUTF8(filename, "<stdin>")) &&
            set_main_loader(dict, filename, "SourceFileLoader") < 0) {
            fprintf(stderr, "python: failed to set __main__.__loader__\n");
            ret = -1;
            goto done;
        }
        v = pyrun_file(fp, filename, Py_file_input, dict, dict,
                       closeit, flags);
    }
    ~ 이하 생략 ~ 
}

pyc file은 이미 컴파일된 python byte code 라 tokenizing, AST step을 생략하다 보니 더 빠르게 실행됩니다.

기본적으로 처음 실행된다 가정하에 pyrun_file method를 조금 더 살펴보면, PyArena 메모리 포인터를 초기화하고 _PyParser_ASTFromFile에서 코드를 AST로 파싱하고 run_mod에서 실행하는 흐름을 확인 할 수 있습니다. run_mode 이후 부터는 아래에서 조금 더 자세히 다루겠습니다.

static PyObject *
pyrun_file(FILE *fp, PyObject *filename, int start, PyObject *globals,
           PyObject *locals, int closeit, PyCompilerFlags *flags)
{
    PyArena *arena = _PyArena_New();
    if (arena == NULL) {
        return NULL;
    }

    mod_ty mod;
    mod = _PyParser_ASTFromFile(fp, filename, NULL, start, NULL, NULL,
                                flags, NULL, arena);

    if (closeit) {
        fclose(fp);
    }

    PyObject *ret;
    if (mod != NULL) {
        ret = run_mod(mod, filename, globals, locals, flags, arena, NULL, 0);
    }
    else {
        ret = NULL;
    }
    _PyArena_Free(arena);

    return ret;
}

PyArena는 single location 기반의 memory pool 이며 프로그램 종료시 PyArena_Free 를 통해 한번에 메모리가 해지되게 됩니다. 참고로 GC 구현은 이와 별개로 더 자세한 내용은 아래 코드를 참고해주세요.
https://github.com/python/cpython/blob/main/Python/gc.c

Tokenizer

_PyParser_ASTFromFile method를 타고 들어가면 parser와 tokenizer를 사용하는 부분을 볼 수 있는데, tokenizer는 코드 블락을 Grammar를 기반으로 tokenizing하고 parser는 이 token들을 사용하여 AST를 build 하는 역할로 이해하면 됩니다.

mod_ty
_PyPegen_run_parser_from_file_pointer(FILE *fp, int start_rule, PyObject *filename_ob,
                             const char *enc, const char *ps1, const char *ps2,
                             PyCompilerFlags *flags, int *errcode,
                             PyObject **interactive_src, PyArena *arena)
{
    struct tok_state *tok = _PyTokenizer_FromFile(fp, enc, ps1, ps2);
    if (tok == NULL) {
        if (PyErr_Occurred()) {
            _PyPegen_raise_tokenizer_init_error(filename_ob);
            return NULL;
        }
        return NULL;
    }
    if (!tok->fp || ps1 != NULL || ps2 != NULL ||
        PyUnicode_CompareWithASCIIString(filename_ob, "<stdin>") == 0) {
        tok->fp_interactive = 1;
    }
    // This transfers the ownership to the tokenizer
    tok->filename = Py_NewRef(filename_ob);

    // From here on we need to clean up even if there's an error
    mod_ty result = NULL;

    int parser_flags = compute_parser_flags(flags);
    Parser *p = _PyPegen_Parser_New(tok, start_rule, parser_flags, PY_MINOR_VERSION,
                                    errcode, arena);
    if (p == NULL) {
        goto error;
    }

    result = _PyPegen_run_parser(p);

AST

AST란? 

추상 구문 트리 - 위키백과, 우리 모두의 백과사전

The abstract syntax tree (AST) is a high-level representation of the program structure without the necessity of containing the source code; it can be thought of as an abstract representation of the source code. The specification of the AST nodes is specified using the Zephyr Abstract Syntax Definition Language (ASDL) [^1], [^2].

코드를 ASDL(Abstract Syntax Definiton Language)라는 high level 언어로 추상화하여 표현한 것입니다. AST의 각 노드들은 statement, experssion 외 몇 가지 특화된 타입으로 구성되며 자세한 구성 요소에 대해선 아래 문서를 참고해주세요.

https://github.com/python/cpython/blob/main/Parser/Python.asdl

cpython/Parser/Python.asdl at main · python/cpython

https://greentreesnakes.readthedocs.io/en/latest/nodes.html

Python에서도 built-in module ast을 사용하여 source code를 AST로 파싱해볼 수 있어 몇 가지 예시를 가져왔습니다.

import ast
import dis


if __name__ == "__main__":
    # Parse source code into AST
    source_code = "a = 1 + 2"
    tree = ast.parse(source_code)
    # For example, let's print the AST
    print(ast.dump(tree))

ast.dump를 호출하게 되면 아래와 같이 Tree를 이루는 객체들이 출력되는 것을 확인할 수 있습니다.

1. Expression, Statement

a = 1 + 2

Module(body=[Assign(targets=[Name(id='a', ctx=Store())], value=BinOp(left=Constant(value=1), op=Add(), right=Constant(value=2)))], type_ignores=[])

2. FunctionDef

@dec1
@dec2
def f(a: 'annotation', b=1, c=2, *d, e, f=3, **g) -> 'return annotation':
  pass

Module(body=[
 FunctionDef(name='f', args=arguments(posonlyargs=[],
   args=[
     arg(arg='a', annotation=Str(s='annotation')),
     arg(arg='b', annotation=None),
     arg(arg='c', annotation=None),
   ], vararg=arg(arg='d', annotation=None), kwonlyargs=[
     arg(arg='e', annotation=None),
     arg(arg='f', annotation=None),
   ], kw_defaults=[
     None,
     Num(n=3),
   ], kwarg=arg(arg='g', annotation=None), defaults=[
     Num(n=1),
     Num(n=2),
   ]), body=[
     Pass(),
   ], decorator_list=[
     Name(id='dec1', ctx=Load()),
     Name(id='dec2', ctx=Load()),
   ], returns=Str(s='return annotation')),
])

지금까지의 과정을 아래 문서에서 다음과 같이 요약설명 하고 있습니다.

https://devguide.python.org/internals/parser/index.html#abstract

몇 가지 검사 후, Parser/parser.c의 헬퍼 함수가 입력받은 소스 코드에 대해 생성 규칙을 적용하기 시작합니다. 이 과정에서 소스 코드를 토큰으로 변환하고, 이 토큰들을 해당하는 규칙과 재귀적으로 매칭합니다. 매칭이 될 때마다 해당 생성 규칙의 규칙 함수가 호출됩니다. 이런 규칙 함수들은 xx_rule 형식을 따릅니다. 여기서 xx는 함수가 처리하는 문법 규칙을 나타내며, 규칙 함수는 Grammar/python.gram 기반 Tools/peg_generator/pegen/c_generator.py에 의해 자동으로 생성됩니다.

각 규칙 함수는 진행하면서 AST 노드를 생성합니다. 이는 필요한 모든 새 노드를 할당하고, 필요한 지원 함수에 대해 적절한 AST 노드 생성 함수를 호출하며, 이들을 필요에 따라 연결함으로써 이루어집니다. 이 과정은 모든 nonterminal 기호가 terminals 기호로 대체될 때까지 계속됩니다. 오류가 발생하면 규칙 함수들은 백트래킹하여 다른 규칙 함수를 시도합니다. 더 이상 적용할 규칙이 없으면 오류가 설정되고 파싱이 종료됩니다.

아래는 위에서 언급된 python.gram에서 생성된 규칙함수 예시입니다.
import sys 와 같은 간단한 named import statement 를 담당하는 규칙함수

// This is the production rule (from python.gram) the rule function
   // corresponds to:
   // import_name: 'import' dotted_as_names
   static stmt_ty
   import_name_rule(Parser *p)
   {
       ...
       stmt_ty _res = NULL;
       { // 'import' dotted_as_names
           ...
           Token * _keyword;
           asdl_alias_seq* a;
           // The tokenizing steps.
           if (
               (_keyword = _PyPegen_expect_token(p, 513))  // token='import'
               &&
               (a = dotted_as_names_rule(p))  // dotted_as_names
           )
           {
               ...
               // Generate an AST for the import statement.
               _res = _PyAST_Import ( a , ...);
               ...
               goto done;
           }
           ...
       }
       _res = NULL;
     done:
       ...
       return _res;
   }

PyCodeObject

파싱된 AST는 _PyAST_Compile 함수에서 byte code로의 compile 과정을 거치게 됩니다. 

static PyObject *
run_mod(mod_ty mod, PyObject *filename, PyObject *globals, PyObject *locals,
            PyCompilerFlags *flags, PyArena *arena, PyObject* interactive_src,
            int generate_new_source)
{
    PyThreadState *tstate = _PyThreadState_GET();
    PyObject* interactive_filename = filename;
    ~ 중략 ~ 
    PyCodeObject *co = _PyAST_Compile(mod, interactive_filename, flags, -1, arena);

세부적으로는 다음과 같이 동작합니다.
1. _PySymtable_Build 함수를 통해 symbol table을 생성합니다. 각 Symbol은 AST node type별 함수 symtable_visit_{node type}  에 의해 생성됩니다. 
2. complier_codegen 함수는 symbol table을 사용하여 pseudo instruction sequence를 생성합니다.
3. pseudo instruction sequence 는 _PyCfg_OptimizeCodeUnit 함수에서 최적화 되어 _PyCfg_OptimizedCfgToInstructionSequence에서 다시 pseudo instruction sequence으로 변환됩니다. 해당 과정은 많은 수의 case statement로 이루어진 코드 블락과도 같습니다.
4. 최적화된 pseudo instruction sequence는 _PyAssemble_MakeCodeObject에 의해 최종적으로 byte code로 변환됩니다.

_PyAssemble_MakeCodeObject 함수의 return type이 PyCodeObject 임을 확인할 수 있습니다.

Python에서도 built-in compile 함수에 AST tree와 실행모드를 넘겨주면 PyCodeObject에 high level에서 접근할 수 있는 인터페이스를 사용할 수 있습니다.

# Compile AST into PyCodeObject
code_object = compile(tree, filename="<ast>", mode="exec")

Byte Code

print(dis.dis(code_object.co_code))

          0 LOAD_CONST               0 (0)
          2 STORE_NAME               0 (0)
          4 LOAD_CONST               1 (1)
          6 RETURN_VALUE

PyCodeObject는 실행될 코드 블락의 byte code 정보를 가지고 있고 co_code property로 접근이 가능합니다. 이는 실행가능한 binary이며 exec built-in 함수로 실행할 수 있습니다.

# Execute the PyCodeObject
exec(code_object)

PyFrameObject

최종적으로 컴파일된 byte code는 _PyEval_EvalFrameDefault 함수에서 evalutaion 됩니다. C++ 코드는 매우 복잡하여 python 코드로 간단하게 알아보겠습니다.

PyCodeObject와 _PyEval_EvalFrameDefault 함수 사이에는 PyFrameObject 가 생성됩니다. PyFrameObject는 실행 중인 함수, 지역 변수, 호출 스택 등 코드 블록의 실행 상태에 대한 정보가 포함된 실행 프레임 정보를 가지고 있습니다.

import inspect

def example_function():
    current_frame = inspect.currentframe()

    # Print information about the current frame
    print("Current Frame:", current_frame)
    print("Frame Code Object:", current_frame.f_code)
    print("Frame Local Variables:", current_frame.f_locals)
    print("Frame Global Variables:", current_frame.f_globals)
    print("Frame Back:", current_frame.f_back)

    return current_frame

frame = example_function()

해당 정보들은 python에서는 built-in module inspect를 사용해서 접근할 수 있습니다.

• f_code: 이 프레임에서 실행 중인 코드 객체.
• f_locals: 이 프레임의 로컬 변수들을 나타내는 dictionary.
• f_globals: 이 프레임의 전역 변수들을 나타내는 dictionary.
• f_back: 이전 스택 프레임(호출자 프레임).

초반부에 python은 stack machine 이라고 했는데요, 때문에 아래와 같이 실행흐름 과정대로 call stack을 navigation 할 수 있습니다.

while frame:
    print("Function Name:", frame.f_code.co_name)
    print("Local Variables:", frame.f_locals)
    print("Global Variables:", frame.f_globals)
    print("Built-in Variables:", frame.f_builtins)
    print("Current Line Number:", frame.f_lineno)
    print("Last Instruction Index:", frame.f_lasti)
    print("Back Frame:", frame.f_back)
    frame = frame.f_back

에러시 trackback 에서 출력해주는 stack trace는 위와 같이 call stack을 올라가면서 정보를 파싱하여 출력된 결과물입니다. debugging, profiling에도 활용될 수 있을 거 같네요.

_PyEval_EvalFrameDefault

PyCodeObject의 byte code가 _PyEval_EvalFrame 함수에 의해 실행되는 자세한 과정은 아래 문서에서 잘설명하고 있습니다.
https://devguide.python.org/internals/interpreter/

위 문서를 탐독하고 나면 bytecode interpreter는 _PyEval_EvalFrameDefault 함수를 가리킨다고 봐도 무방함을 이해할 수 있습니다.

PEP 523 – Adding a frame evaluation API to CPython

https://peps.python.org/pep-0523/

PEP 523 – Adding a frame evaluation API to CPython | peps.python.org

그 중 저는 PEP523에서 제안된 custom interpreter 함수를 사용할 수 있게 해주는 CPython API 가 흥미로웠습니다.

간단히 말해서 커스터마이징한 _PyEval_EvalFrame 함수를 사용할 수 있게 된 것인데요, 이는 원래 third party에게 JIT Interpreter를 구현할 인터페이스를 제공하기 위해 고안되었다고 합니다.

Pytorch framework을 사용하는 분들은 익숙하신 TorchDynamo tracer 구현에서도 위 CPython API를 내부적으로 사용하는데요,
tracer는 모델의 forward 함수가 실행될 때 해당 bytecode를 실행직전 수정하는 custom interpreter를 사용합니다. bytecode를 분석하여 tensor 연산에 사용되는 모든 operator를 graph로 표현 및 PyFrameObject로 저장하게 됩니다.
(해당 정보를 별도 디자인한 클래스로 저장하긴함).
그림으로 표현하면 아래와 같습니다.

https://pytorch.org/docs/stable/torch.compiler_dynamo_overview.html

tracer는 보통 각 하드웨어별 최적화된 operator로 변환해주는 compiler와 함께 사용되어 최종적으로 아래와 같은 그림이 되게 됩니다.
그리고 위 그림에서 보이는 Guards 라는 개념이 등장하는데, 이는 forward의 input tensor의 아래 Property를 확인하여

• Python class of the tensor (tensor subclassing, etc)
• dtype
• device
• requires_grad
• dispatch_key (with thread-local includes/excludes applied)
• ndim
• sizes*
• strides*

동일하지 않다고 판단되면 오른쪽 분기를 태워 recompile을 하여 실행하는 방식입니다.

위는 torch.compile (JIT complier) 기준 동작 방식이며 torch_tensorrt 와 같은 AOT compiler는 recompile을 수행하지 않고 에러를 발생하는 식입니다. 
참고 (https://surgach.tistory.com/142)

추론 최적화 시리즈 [1] Bert4rec Pytorch module을 Torch-Tensorrt로 compile 하여 Tritonserver로 실시간 추론하

더 자세한 내용은 아래 두 문서를 참고해주세요.
- https://pytorch.org/docs/stable/torch.compiler_dynamo_overview.html

Dynamo Overview — PyTorch 2.4 documentation

- https://pytorch.org/docs/main/torch.compiler_dynamo_deepdive.html

Dynamo Deep-Dive — PyTorch main documentation

 

Summary

1. Python file은 아래와 같은 메커니즘을 통해 실행되는 것을 코드레벨에서 이해해 보았습니다.
: Python Code -> tokenizing -> AST -> python byte code(pyc) -> PyCodeObject -> PyFrameObject -> PyFrameEx

2. CPython의 전반적인 아키텍처에 대해서 알아보았습니다.
3. PEP523 CPython API 를 활용한 TorchDynamo tracer에 대해 알아보았습니다.

세부적인 동작들에 대해선 더 공부해야할 부분들이 많이 남아있지만, CPython의 핵심 구현을 이해할 수 있었고 Python에 조금 더 애착이 생기게 된 거 같습니다 ㅎㅎ :) 

긴 글 읽어주셔서 감사합니다.

 

Reference

https://devguide.python.org/internals/interpreter/

https://peps.python.org/pep-0523/

PEP 523 – Adding a frame evaluation API to CPython | peps.python.org

https://docs.python.org/3/c-api/veryhigh.html#c.PyEval_EvalFrameEx

The Very High Level Layer

https://github.com/python/cpython

GitHub - python/cpython: The Python programming language

https://pytorch.org/docs/stable/torch.compiler_dynamo_overview.html

Dynamo Overview — PyTorch 2.4 documentation

https://pytorch.org/docs/main/torch.compiler_dynamo_deepdive.html

Dynamo Deep-Dive — PyTorch main documentation