OneFlow 源码阅读 12：从 Tensor 看 CPython 的对象创建过程

春节前后拜读了许啸宇的《TorchDynamo初探：Python ByteCode的动态修改》，随后简单总结了一下 TorchDynamo 的执行过程。这期间对 CPython 也多了一些了解。记得之前看到过，oneflow Tensor 是通过 CPython API 导出给 Python 环境的，当时很多细节都不了解，就着这个热乎劲，重新看了一下 Tensor 类型注册相关的内容，顺便熟悉一下 CPython 创建对象的过程。
本文中涉及 CPython 相关的内容可以参考 Python behind the scenes。

1 Python 中的类型

Python 中一切都是对象。对象的数据类型本身也是对象，比如 float、list 这些，都是对象。这些类型对象中，有一个比较特殊，就是 type，它是元类型（metatype），即类型的类型。以下断言是成立的：

assert(float.__class__ is type)
assert(type.__class__  is type)

这些类型对象在 CPyhton 中的定义如下，它们的 C 类型都是 PyTypeObject，也都可以被视为 PyObject。

• object: PyBaseObject_Type（定义）
• type: PyType_Type（定义）
• float: PyFloat_Type（定义）
• list: PyList_Type（定义）

PyTypeObject 的很多字段（slot）都是函数指针。这些 slot 的不同取值，决定了不同类型的行为。slotdefs 数组定义了 magic method 的 slot。

2 Tensor 类型的注册

one::Tensor 类型注册的代码在 tensor.cpp 中。这里定义了三个 PyTypeObject 对象：

• TensorMetaclass_Type: metatype
• PyTensorObject_Type: oneflow.Tensor
• PyParameterObject_Type: oneflow.Parameter，是 oneflow.Tensor 的子类型。

MakeTensorMetaclass 函数负责初始化 TensorMetaclass_Type 对象。下面这行代码创建对象实体：

auto* heap_type = (PyHeapTypeObject*)PyType_Type.tp_alloc(&PyType_Type, 0);

这行代码有两点需要澄清：

• PyType_Type.tp_alloc 是啥？
• 返回的类型为何是 PyHeapTypeObject？

2.2 CPython 的类型初始化

从字面意义看，tp_alloc 应该会创建一个对象并返回其指针。但是在定义 PyType_Type 时，tp_alloc 的默认值是 0。所有类型在初始化时都会从 base type 拷贝一些 slot。PyType_Type.tp_alloc 的值是从 PyBaseObject_Type.tp_alloc 拷贝而来的，具体函数是 PyType_GenericAlloc（可以参考 Python behind the scenes #6: how Python object system works 中的 Slot inheritance 一节）。

CPython 在启动时初始化类型的具体过程如下：
pycore_init_types -> _PyTypes_Init -> INIT_TYPE(PyType_Type) -> PyType_Ready -> type_ready(...)

• type_ready_set_bases(...): 设置type->tp_base = &PyBaseObject_Type，如果 tp_base 为空。

• type_ready_inherit

  • inherit_slots: 拷贝 slots

    • COPYSLOT(tp_alloc): 拷贝单个 slot

2.3 tp_alloc 返回的类型为何是 PyHeapTypeObject？

PyType_GenericAlloc 声明的返回的类型是 PyObject *。不过在计算对象的 size 时，不会小于 tp_basicsize，PyType_Type.tp_basicsize 被初始化为 sizeof(PyHeapTypeObject)，所以这里返回的对象可以放心地作为 PyHeapTypeObject 使用。CPython 的 PyType_FromModuleAndSpec 中也是这样处理的。

2.3.1 PyHeapTypeObject 支持更灵活的动态类型

CPython 的内置类型都是直接通过 PyTypeObject 的静态变量定义的，也就是所谓的 Static Types。这些类型对象的 slot 值都可以在编译期获得。

不同于静态定义的 PyTypeObject，PyHeapTypeObject 可以在运行时动态定制类型的行为，比如在框架启动后根据配置文件、外部输入等信息确定类型的 slot 的值。这就是所谓的 Heap Types，类型对象是在堆中存储的。

在 Python 下执行 my_type = type('my_type', (), {}) 会进入 type_new_alloc 函数，这里会将 PyHeapTypeObject 的对象类型的 slot 字段地址赋值给指针类型的 slot 字段。然后用户再对 PyHeapTypeObject 的字段赋值，就可以根据外部输入创建动态类型了。直接调用 PyType_Type.tp_call(&PyType_Type, ...) 应该可以达到类似效果。

2.4 tp_flags 值的含义

Tensor 类型和元类型设置了三个标签：

• Py_TPFLAGS_DEFAULT: 这应该是所有对象都会设置的一个标签。
• Py_TPFLAGS_BASETYPE: 表示可以作为其它类型的父类。
• Py_TPFLAGS_HEAPTYPE: 从堆中分配的 type object 需要设置这个标记。创建这种类型的对象时，会递增 type object 的引用计数；释放时递减 type object 的引用计数。

2.5 TensorMetaclass_Type 的主要 slot 值

• ob_type: &PyType_Type
• tp_name: _TensorMeta
• tp_base: &PyType_Type
• tp_call: TensorMetaCls_call
• tp_dealloc: TensorMetaCls_dealloc
• tp_alloc: PyType_GenericAlloc，继承自 PyBaseObject_Type。

2.6 PyTensorObject_Type 的主要 slot 值

Tensor 的类型对象 PyTensorObject_Type 是通过元类 TensorMetaclass_Type 创建的。PyTensorObject_Type 的部分字段如下：

• ob_type: &TensorMetaclass_Type
• tp_name: Tensor
• tp_init: PyTensorObject_init，调用 functional::_legacy_tensor_ctor(...)
• tp_dealloc: PyTensorObject_dealloc
• tp_alloc: PyType_GenericAlloc，继承自 PyBaseObject_Type。

3 Tensor 对象的创建过程

echo "Tensor()" | python3 -m dis 显示的 Python 子节码如下：

  1           0 LOAD_NAME                0 (Tensor)
              2 CALL_FUNCTION            0
              4 POP_TOP
              6 LOAD_CONST               0 (None)
              8 RETURN_VALUE

LOAD_NAME 根据名字 Tensor 加载对象，也就是把 &PyTensorObject_Type 放到 Value stack 的栈顶。

CALL_FUNCTION 指令会调用 call_function 函数。这个函数的主要代码如下：

    PyObject **pfunc = (*pp_stack) - oparg - 1;
    PyObject *func = *pfunc;
    PyObject *x, *w;
    Py_ssize_t nkwargs = (kwnames == NULL) ? 0 : PyTuple_GET_SIZE(kwnames);
    Py_ssize_t nargs = oparg - nkwargs;
    PyObject **stack = (*pp_stack) - nargs - nkwargs;

    if (trace_info->cframe.use_tracing) {
        x = trace_call_function(tstate, trace_info, func, stack, nargs, kwnames);
    }
    else {
        x = PyObject_Vectorcall(func, stack, nargs | PY_VECTORCALL_ARGUMENTS_OFFSET, kwnames);
    }

oparg 是 opcode 对应的参数，在这个例子中就是 LOAD_NAME 对应的 0。

pp_stack 是 value stack 的栈顶指针。所以 func 就是栈顶元素，也就是 &PyTensorObject_Type。

PyObject_Vectorcall 会调用 _PyObject_VectorcallTstate 继续处理。_PyObject_VectorcallTstate 中的核心代码如下：

func = PyVectorcall_Function(callable);
if (func == NULL) {
    Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
    return _PyObject_MakeTpCall(tstate, callable, args, nargs, kwnames);
}
res = func(callable, args, nargsf, kwnames);

这里callable 就是 &PyTensorObject_Type。因为 TensorMetaclass_Type 没有设置 Py_TPFLAGS_HAVE_VECTORCALL，会直接返回 NULL。然后继续调用 _PyObject_MakeTpCall。这个函数的核心代码如下：

ternaryfunc call = Py_TYPE(callable)->tp_call;
PyObject *result = NULL;
result = call(callable, argstuple, kwdict);

这里的 call 是 TensorMetaclass_Type.tp_call，也就是 TensorMetaCls_call。所以 Tensor 对象的创建是由 TensorMetaCls_call 完成的。这个函数直接转发给 PyType_Type.tp_call，也就是 type_call 函数，其中的 type 参数就是 &PyTensorObject_Type。

3.1 其它设置 Py_TPFLAGS_HAVE_VECTORCALL 的类型

class A:
    def func(self):
        print("in func")

a = A()
print(a.func.__class__)
print(A.func.__class__)

除了 PyType_Type，以下内置类型也设置了 Py_TPFLAGS_HAVE_VECTORCALL，对类似 f() 这种操作，它们都有自己定制的处理逻辑。

• method_descriptor: 比如 list.append
• function: 比如 A.func。
• method: 比如 a.func。所以，如果是函数或方法，PyVectorcall_Function 返回的就不是空指针，会直接调用 func。
• builtin_function_or_method，比如 sum。

3.2 type_call 的执行

type_call 的主要流程如下：

• obj = type->tp_new(type, args, kwds)，tp_new 的值是从 PyBaseObject_Type 继承来的 object_new。所以实际执行的是 object_new(&PyTensorObject_Type, ...)

  • type->tp_alloc(type, 0)。实际调用继承自 PyBaseObject_Type 的 PyType_GenericAlloc，即 PyType_GenericAlloc(&PyTensorObject_Type, 0)。

    • size 是需要分配的内存大小。至少是 tp_basicsize，即 sizeof(PyTensorObject)。
    • PyObject_Malloc
    • memset(obj, '\0', size)

    • _PyObject_Init

      • 设置对象类型
      • 增加 PyTensorObject_Type 的引用计数。
      • 调用 _Py_NewReference 将新 Tensor 对象的引用计数设置为 1。

• type->tp_init(obj, args, kwds)。type 是 &PyTensorObject_Type，实际调用的是 PyTensorObject_init。

  • functional::_legacy_tensor_ctor
  • PyTensor_wrap，self 是要初始化的对象。

functional::_legacy_tensor_ctor 返回的已经是一个 PyTensorObject 指针了，估计可能是原来适配 pybind11 的缘故。但是现在数据要绑定到 self 上，所以需要从 temp 拷贝切换到 self。

3.3 heap 对象影响其 type object 的引用计数

每创建一个 flow.Tensor 对象，PyTensorObject_Type 的引用计数会递增，只要不主动删除，类型对象不会自动释放。

import sys
import oneflow as flow
print(sys.getrefcount(flow.Tensor))
a = flow.ones(2, 3)
print(sys.getrefcount(flow.Tensor))
del a
print(sys.getrefcount(flow.Tensor))

通过 PyType_GenericAlloc 创建对象时，如果是 heap type，会递增 type object 的引用计数。对象释放时，也需要递减 type object 的引用计数。

3.4 普通内置类型的创建

Python 内置的 list、float 这些的 ob_type 是 &PyType_Type，它设置了 Py_TPFLAGS_HAVE_VECTORCALL，初始化 tp_vectorcall 为 type_vectorcall，初始化 tp_vectorcall_offset 为 tp_vectorcall 在 struct 中的偏移。所以 PyVectorcall_Function 返回的就是 type_vectorcall，这些内置类型就用这个函数构造对象。以 float 为例，_PyObject_VectorcallTstate 中实际调用的是：

res = type_vectorcall(&PyFloat_Type, args, nargsf, kwnames);

type_vectorcall 中的 metatype、_PyObject_MakeTpCall 中的 callable 是 &PyFloat_Type。最终仍是通过 type_call 完成对象创建。

4 PyTypeOjbect 几个 slot 的语义

• tp_call: 对于自定义类型来说，如果 A 是一个 Python 类型，执行 A()，最终会调用它的元类型的 tp_call，即 A->ob_type->tp_call()。

  • 这一点和 C++ 是不一样的。可以这样理解：C++ 中，编译器在栈上为对象分配存储，或者通过调用 new 在堆上分配存储；然后调用类的构造函数初始化对象。而 CPython 中，虚拟机调用 metatype 的 tp_call 统一完成存储分配和对象初始化，其中对象初始化由类型 A 完成。
  • PyBaseObject_Type.tp_call 是 0，因为不是每个对象都是可调用的。

• tp_dealloc: Python 运行时释放一个对象时（例如 oneflow.Tensor），会调用该类型的 tp_dealloc slot（例如 PyTensorObject_dealloc）。调用这个 slot 时，对象实例 self 仍存在，但是引用计数已经为 0。这个函数需要释放 self 持有的所有其它对象的引用，以及其它相关资源；然后调用类型的 tp_free slot，即 self->ob_type->tp_free(self)。如果是 heap 类型，还需要递减 type object 的引用计数。

  • PyBaseObject_Type.tp_dealloc 是 object_dealloc
  • PyType_Type.tp_dealloc 是 type_dealloc

• tp_init: 完成对象的初始化，相当于 class 的 __init__ 方法。

  • PyBaseObject_Type.tp_init 是 object_init
  • PyType_Type.tp_init 是 type_init
• tp_alloc: 为类型的实例分配存储空间。创建元类型时调用的是 PyType_Type.tp_alloc，因为 _TensorMeta、float、list 等的 ob_type 都是 PyType_Type。创建 Tensor 类型时调用的是 TensorMetaclass_Type.tp_alloc；创建 Tensor 对象时调用的是 PyTensorObject_Type.tp_alloc。

• tp_new: 创建对象实例。必须调用 tp_alloc 为对象分配存储空间。tp_new 应尽量少做初始化工作；初始化尽可能在 tp_init 中进行。

  • PyBaseObject_Type.tp_new 是 object_new
  • PyType_Type.tp_new 是 type_new

• tp_free: 释放 Python 对象存储。

  • tp_dealloc 中需要调用 tp_free，例如 PyTensorObject_dealloc。
  • PyBaseObject_Type.tp_free 是 PyObject_Del
  • PyType_Type.tp_free 是 PyObject_GC_Del

5 metatype 的作用是什么？

目前看，TensorMetaclass_Type 的主要作用，应该是在 CPython 层面对创建 PyTensorObject 对象的行为进行定制。

元类的作用是创建对象，并进行 Python 层面的初始化，比如设置 ob_type、引用计数等。tensor type 的作用，是进行特定类型的、业务相关的初始化。例如 Tensor 和 Parameter 有各自的初始化函数。

如果没有 TensorMetaclass_Type，就会像 float、list 那样，调用 PyType_Type 的 vectorcall，最终仍会调用 type_call。

6 避免 PyTensorObject 被频繁创建、释放

在 v0.9.0 中，Tensor 对象初始化时，_legacy_tensor_ctor 已经返回了 PyTensorObject 对象的指针。为什么还要用 PyTensor_wrap 处理一下呢？

在 v0.8.0 中，以下代码中输出的 id(a.grad) 的值是不固定的：

# test.py
import oneflow as flow

a = flow.ones(2, 3).requires_grad_()
b = flow.ones(2, 3).requires_grad_()

c = a + b
s = c.sum()
s.backward()

print(id(a.grad))
print(id(a.grad))

oneflow 的 issue 9500 对此作了详细讨论。简单的说，这种现象是由如下原因造成的：

• backward 返回后，grad 数据只是以 one::Tensor 的形态存在，从未暴露给 Python 端，其指向 PyTensorObject 的 pyobject_ 指针还是空的。
• Python 环境下 id(a.grad) 的调用会创建一个新的 PyTensorObject 对象，但是 a.grad 是一个临时对象，CPython 的 call_function 返回前需要调用 Py_DECREF 递减引用计数。
• 因为是临时对象、没有其它引用，调用 Py_DECREF 使引用计数为 0、导致 PyTensorObject_dealloc 被调用，pyobject_ 指针重新变为空。但是在 C++ 层面，grad 并没有被析构，还在 Tensor a 中保存着。
• 这样每次调用 id(a.grad) 都会重复创建、释放 PyTensorObject 对象，打印的 id 也就不是固定的。（但是对 one::Tensor 对象没有影响）

pull 9544 中已经用一种巧妙的方案修复这个问题。

6.1 C++ 对象与 Python 对象的关系

C++ 是后端，Python 是前端。C++ 中的对象需要封装为 PyOjbect，才能在 Python 中供用户使用。

用一个比喻来形容二者的关系：C++ 对象是真身，PyObject 是幻影分身。真身是基础；幻影分身无法脱离真身而独立存在。也就是说，可以只有 C++对象、而没有 PyObject；但是 PyObject 不能在没有 C++ 对象的情况下独立存在。

6.2 Tensor 对象的状态

之所以会出现 issue 9500 中的问题，原因是之前认为对象只有两种状态：

• 未返回给 Python 端，或者引用计数已经为 0。
• 返回给 Python 端，引用计数不为 0。

如果只允许这两种状态，就会出现上述频繁创建对象的情况。所以 pull 9544 引入了一个新的 bool 变量，和原来的指针一起表示对象的状态。

• pyobj_ptr_: PyTensorObject 的智能指针。deleter 函数用于在 one::Tensor 对象析构时处理 PyTensorObject。
• owns_pyobj_: bool 变量，表示该 one::Tensor 是否持有对 PyTensorObject 的管理权限。

这样，one::Tensor 在生存期内就有三种合法状态：

• 状态 0：初始状态。one::Tensor 对象刚在 C++ 中创建完毕，尚未返回给 Python 端。grad 在初始时就是这种状态。

  • owns_pyobj_: false，没有 PyTensorObject 对象需要管理。
  • pyobj_ptr_: nullptr。

• 状态 1：对象返回给 Python 端，引用计数大于 0。a = flow.ones(2, 3) 就是这种状态。

  • owns_pyobj_: false，由 Python 运行时管理 PyTensorObject 对象。
  • pyobj_ptr_: 不空。

• 状态 2：对象返回给 Python 端，但后来引用计数变为 0。a = flow.ones(2, 3); del a 以及 id(a.grad) 就是这种状态。

  • owns_pyobj_: true，由 one::Tensor 管理 PyTensorObject 对象。
  • pyobj_ptr_: 持有 &PyTensorObject，引用计数为1。

pyobj_ptr_ == nullptr && owns_pyobj_ == true 是一种非法状态。

6.3 Tensor 对象状态的转换

one::Tensor 的三种状态之间的转换关系如下：

在 v0.9.0 的代码中，PyTensor_tryResurrect 的作用是从状态 1 转为状态 2。之后 one::Tensor 可能会被析构，也可能不会。

PyTensor_wrap 的作用是从各种合法状态转为状态 2。三种状态转换的代码对应关系如下：

• 0 -> 1
• 1 -> 1
• 2 -> 1

6.4 PyTensor_tryResurrect 何时返回 false？

PyTensor_tryResurrect 中，只有在 one::Tensor 析构时，self->data 才是空的。调用链路如下：

• 进入 PyTensor_tryResurrect 时，如果 Python 环境中只有 self 一个引用，会将 PyTensorObject.data 置为空，这时只有 tensor 一个指针持有 one::Tensor 对象。
• PyTensor_tryResurrect 返回前，tensor 智能指针析构，依次导致 one::Tensor 及其成员变量 pyobj_ptr_ 析构。
• pyobj_ptr_ 析构执行 deleter 函数，递减 PyTensorObject 的引用计数为 0，导致再次递归进入 PyTensorObject_dealloc，并递归进入 PyTensor_tryResurrect。
• 此时在 PyTensor_tryResurrect 中，PyTensorObject.data 为空，返回false。
• 递归调用的 PyTensorObject_dealloc 中执行资源清理操作。
• 首次调用的 PyTensorObject_dealloc 直接返回。

7 附录

7.1 oneflow gdb 断点示例

启动 gdb

source /mnt/oneflow/build/source.sh
gdb --args python3

断点示例

set breakpoint pending on
# run 之前设置
break oneflow::one::TensorMetaCls_call

# flow.Tensor() 之后设置
break PyType_GenericAlloc
break typeobject.c:1169
break oneflow::one::PyTensorObject_init

# issue 9500 的断点
# backward 之前设置
break oneflow::one::Tensor::set_pyobject
# backward 之后设置
break oneflow::one::PyTensor_New
break oneflow::one::Tensor::~Tensor

7.2 CPython 宏展开命令示例

gcc -E -std=c99 \
    -DCONFIG_64 \
    -DPy_BUILD_CORE \
    -D_POSIX_THREADS \
    -I Include \
    -I Include/internal \
    -I Modules/_decimal/libmpdec \
    -I PC \
    Objects/typeobject.c > typeobject.c

8 参考资料

• Python behind the scenes（中文）

  • #3: stepping through the CPython source code
  • #4: how Python bytecode is executed
  • #6: how Python object system works

• Python 官方文档

  • Data model
  • Defining Extension Types: Tutorial
  • Type Objects

• PyTensorObject 频繁创建、释放问题的修复

  • github issue
  • pull request
• CPython v3.10.9
• oneflow v0.9.0